(*
 * Copyright 2008 Savonet team
 *
 * This file is part of ocaml-speex.
 *
 * ocaml-speex is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * ocaml-speex is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with ocaml-speex; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

(**
  * An wav to ogg converter using ocaml-speex.
  *
  * @author Samuel Mimram, Romain Beauxis, and many others...
  *)

open Speex
open Unix


let src = ref ""
let dst = ref ""

let input_string chan len =
  let ans = String.create len in
    really_input chan ans 0 len ;
    ans

let input_int chan =
  let buf = input_string chan 4 in
    (int_of_char buf.[0])
    + (int_of_char buf.[1]) lsl 8
    + (int_of_char buf.[2]) lsl 16
    + (int_of_char buf.[3]) lsl 24

let input_short chan =
  let buf = input_string chan 2 in
    (int_of_char buf.[0]) + (int_of_char buf.[1]) lsl 8

let usage = "usage: wav2speex [options] source destination"
let float = ref false
let mode =  ref Wideband
let fpp = ref 5
let vbr = ref false
let quality = ref 7

let _ =
  let f = Printf.sprintf in
  let string_of_mode m = 
    match m with
      | Narrowband     -> "narrowband"
      | Wideband       -> "wideband"
      | Ultra_wideband -> "ultra-wideband"
  in
  Arg.parse
    [
      "--float", Arg.Unit (fun b -> float := true ),
      f "Use floats for decoding. Default: %b" !float;
      "--mode", Arg.String 
        (fun b -> 
          match b with
            | "narrowband"     -> mode := Narrowband
            | "wideband"       -> mode := Wideband
            | "ultra-wideband" -> mode := Ultra_wideband
            | _ -> failwith "unkown mode"),
      f "Encoding mode, one of \"narrowband\", \"wideband\" or \"ultra-wideband\". \
         Default: %s" (string_of_mode !mode);
      "--frame_per_packet", Arg.Int (fun b -> fpp := b ),
      f "Frames per Ogg packet. Default: %i" !fpp;
      "--vbr", Arg.Unit (fun b -> vbr := true ),
      f "Encode in vbr mode. Default: %b" !vbr;
      "--quality", Arg.Int (fun b -> quality := b ),
      f "Encoding bitrate, in Kbps. Default: %i" !quality;
    ]
    (
      let pnum = ref (-1) in
        (fun s -> incr pnum; match !pnum with
           | 0 -> src := s
           | 1 -> dst := s
           | _ -> Printf.eprintf "Error: too many arguments\n"; exit 1
        )
    ) usage;
  if !src = "" || !dst = "" then
    (
      Printf.printf "%s\n" usage;
      exit 1
    );
  let ic = open_in_bin !src in
  let oc = open_out_bin !dst in
    (* TODO: improve! *)
    if input_string ic 4 <> "RIFF" then invalid_arg "No RIFF tag";
    ignore (input_string ic 4);
    if input_string ic 4 <> "WAVE" then invalid_arg "No WAVE tag";
    if input_string ic 4 <> "fmt " then invalid_arg "No fmt tag";
    let _ = input_int ic in
    let _ = input_short ic in (* TODO: should be 1 *)
    let channels = input_short ic in
    let infreq = input_int ic in
    let _ = input_int ic in (* bytes / s *)
    let _ = input_short ic in (* block align *)
    let ibits = input_short ic in
    if input_string ic 4 <> "data" then invalid_arg "No data tag";
    let _ = input_int ic in (* datalen *)
    let fos buf =
      let len = String.length buf / (2 * channels) in
      let ans = Array.init channels (fun _ -> Array.create len 0) in
        for i = 0 to len - 1 do
          for c = 0 to channels - 1 do
            let n =
              int_of_char buf.[2 * channels * i + 2 * c]
              + int_of_char buf.[2 * channels * i + 2 * c + 1] lsl 8
            in
            let n =
              if n land 1 lsl 15 = 0 then
                n
              else
                (n land 0b111111111111111) - 32768
            in
             ans.(c).(i) <- max (-32768) (min 32767 n);
          done;
        done;
        ans
    in
    let enc = Encoder.init !mode !fpp in
    if not !vbr then
      Encoder.set enc SPEEX_SET_QUALITY !quality
    else
      Encoder.set enc SPEEX_SET_VBR_QUALITY !quality;
    Encoder.set enc SPEEX_SET_SAMPLING_RATE infreq;
    let ivbr = 
      if !vbr then 1 else 0 
    in
    Encoder.set enc SPEEX_SET_VBR ivbr;
    let fsize = Encoder.get enc SPEEX_GET_FRAME_SIZE in
    let os = Ogg.Stream.create () in
    let header = Header.init ~rate:infreq ~nb_channels:channels 
                             ~mode:(!mode) ~vbr:(!vbr) 
                             ~frames_per_packet:(!fpp) 
                             () in
    Header.encode_header header [] os;
    output_string oc (Ogg.Stream.flush os);
    let start = Unix.time () in
    let smode = 
      match !mode with
        | Narrowband     -> "narrowband"
        | Wideband       -> "wideband"
        | Ultra_wideband -> "ultra-wideband"
    in
      Printf.printf
        "Input detected: PCM WAVE %d channels, %d Hz, %d bits\n%!"
        channels infreq ibits;
      Printf.printf
        "Encoding to: SPEEX %d channels, %d Hz, %s, VBR: %s\nPlease wait...\n%!"
        channels infreq smode (string_of_bool !vbr);
        begin try while true do
          let buflen = 2*fsize*channels in
          let buf = String.create buflen in
          let feed () =
            really_input ic buf 0 buflen;
            let fbuf = fos buf in
            assert(Array.length (fbuf.(0)) = fsize);
            fbuf
          in
          let (h,v) =
            if channels = 1 then
              let feed () = 
                let frame = feed () in
                frame.(0)
              in
              if !float then
                let feed () = 
                  Array.map float_of_int (feed ())
                in 
                Encoder.encode_page enc os feed
              else
                Encoder.encode_page_int enc os feed   
            else
              if !float then
                let feed () = 
                  Array.map (fun x -> Array.map float_of_int x) (feed ())
                in
                Encoder.encode_page_stereo enc os feed
              else
                Encoder.encode_page_int_stereo enc os feed
          in
          output_string oc (h ^ v);
        done;
        with End_of_file -> () end;
        Encoder.eos enc os;
        output_string oc (Ogg.Stream.flush os);
        close_in ic;
        close_out oc;
        Printf.printf "Finished in %.0f seconds.\n" ((Unix.time ())-.start);
        Gc.full_major ()