(*
 * Copyright 2003-2011 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
 *)

(**
  * Functions for decoding speex files using libspeex.
  *
  * @author Romain Beauxis
  *)

exception Invalid_frame_size

type mode = Narrowband | Wideband | Ultra_wideband

(* Internal use only *)
type internal_mode

external internal_mode_of_int : int -> internal_mode = "caml_speex_get_mode"

(* Values from speex.h *)
let mode_of_int x =
  match x with
    | 0 -> Narrowband
    | 1 -> Wideband
    | 2 -> Ultra_wideband
    | _ -> failwith "unknown mode"

let int_of_mode x = 
  match x with
    | Narrowband     -> 0
    | Wideband       -> 1
    | Ultra_wideband -> 2

let internal_mode_of_mode x = internal_mode_of_int (int_of_mode x)

(** TODO: better implementation.. *)
(* Generated by control_define *)
type control =
  SPEEX_RESET_STATE |
  SPEEX_SET_ENH |
  SPEEX_GET_ENH |
  SPEEX_GET_FRAME_SIZE |
  SPEEX_SET_QUALITY |
  SPEEX_SET_MODE |
  SPEEX_GET_MODE |
  SPEEX_SET_LOW_MODE |
  SPEEX_GET_LOW_MODE |
  SPEEX_SET_HIGH_MODE |
  SPEEX_GET_HIGH_MODE |
  SPEEX_SET_VBR |
  SPEEX_GET_VBR |
  SPEEX_SET_VBR_QUALITY |
  SPEEX_GET_VBR_QUALITY |
  SPEEX_SET_COMPLEXITY |
  SPEEX_GET_COMPLEXITY |
  SPEEX_SET_BITRATE |
  SPEEX_GET_BITRATE |
  SPEEX_SET_SAMPLING_RATE |
  SPEEX_GET_SAMPLING_RATE |
  SPEEX_SET_VAD |
  SPEEX_GET_VAD |
  SPEEX_SET_ABR |
  SPEEX_GET_ABR |
  SPEEX_SET_DTX |
  SPEEX_GET_DTX |
  SPEEX_SET_SUBMODE_ENCODING |
  SPEEX_GET_SUBMODE_ENCODING |
  SPEEX_SET_PLC_TUNING |
  SPEEX_GET_PLC_TUNING |
  SPEEX_SET_VBR_MAX_BITRATE |
  SPEEX_GET_VBR_MAX_BITRATE |
  SPEEX_SET_HIGHPASS |
  SPEEX_GET_HIGHPASS |
  SPEEX_GET_ACTIVITY

let int_of_control x =
    match x with
      | SPEEX_SET_ENH -> 0
      | SPEEX_GET_ENH -> 1
      | SPEEX_GET_FRAME_SIZE -> 3
      | SPEEX_SET_QUALITY -> 4
      | SPEEX_SET_MODE -> 6
      | SPEEX_GET_MODE -> 7
      | SPEEX_SET_LOW_MODE -> 8
      | SPEEX_GET_LOW_MODE -> 9
      | SPEEX_SET_HIGH_MODE -> 10
      | SPEEX_GET_HIGH_MODE -> 11
      | SPEEX_SET_VBR -> 12
      | SPEEX_GET_VBR -> 13
      | SPEEX_SET_VBR_QUALITY -> 14
      | SPEEX_GET_VBR_QUALITY -> 15
      | SPEEX_SET_COMPLEXITY -> 16
      | SPEEX_GET_COMPLEXITY -> 17
      | SPEEX_SET_BITRATE -> 18
      | SPEEX_GET_BITRATE -> 19
      | SPEEX_SET_SAMPLING_RATE -> 24
      | SPEEX_GET_SAMPLING_RATE -> 25
      | SPEEX_RESET_STATE -> 26
      | SPEEX_SET_VAD -> 30
      | SPEEX_GET_VAD -> 31
      | SPEEX_SET_ABR -> 32
      | SPEEX_GET_ABR -> 33
      | SPEEX_SET_DTX -> 34
      | SPEEX_GET_DTX -> 35
      | SPEEX_SET_SUBMODE_ENCODING -> 36
      | SPEEX_GET_SUBMODE_ENCODING -> 37
      | SPEEX_SET_PLC_TUNING -> 40
      | SPEEX_GET_PLC_TUNING -> 41
      | SPEEX_SET_VBR_MAX_BITRATE -> 42
      | SPEEX_GET_VBR_MAX_BITRATE -> 43
      | SPEEX_SET_HIGHPASS -> 44
      | SPEEX_GET_HIGHPASS -> 45
      | SPEEX_GET_ACTIVITY -> 47


let _ =
  Callback.register "caml_speex_mode_of_int" mode_of_int;
  Callback.register "caml_speex_int_of_mode" int_of_mode;
  Callback.register_exception "ocaml_speex_invfrlen_exn" Invalid_frame_size;
  Callback.register_exception "ocaml_speex_eos_exn" Ogg.End_of_stream

module Header = 
struct
 
  type t = 
    { id:                     string;
      version:                string;
      version_id:             int;
      header_size:            int; 
      rate:                   int;
      mode:                   mode;
      mode_bitstream_version: int;
      nb_channels:            int;
      bitrate:                int;
      frame_size:             int;
      vbr:                    bool;
      frames_per_packet:      int;
      extra_headers:          int
    }

  (* Defined in speex_header.h *)
  let header_string_length = 8
  let header_version_length = 20
 
  external init : int -> int -> internal_mode -> int -> bool -> t = "caml_speex_init_header" 
  
  let init ?(frames_per_packet=1) ?(mode=Wideband) ?(vbr=true) ~nb_channels ~rate () = 
    init rate nb_channels (internal_mode_of_mode mode) frames_per_packet vbr

  external encode_header_packetout : t -> string array -> Ogg.Stream.packet*Ogg.Stream.packet = "caml_speex_encode_header"

  let encode_header_packetout e l = 
    let l = 
      List.map (fun (x,y) -> Printf.sprintf "%s=%s" x y) l
    in
    encode_header_packetout e (Array.of_list l)

  let encode_header e l s =
    let p1,p2 = encode_header_packetout e l in
    Ogg.Stream.put_packet s p1;
    Ogg.Stream.put_packet s p2

  external header_of_packet : Ogg.Stream.packet -> t = "caml_speex_header_of_packet"

  external comments_of_packet : Ogg.Stream.packet -> string array = "caml_speex_comments_of_packet"

  let comments_of_packet p = 
    let x = comments_of_packet p in
    let vendor = x.(0) in
    let x = Array.sub x 1 (Array.length x - 1) in
    let c_k = ref 0 in
    let c_v = ref 0 in
    let split s = 
      try
        let i = String.index s '=' in
        try
          String.sub s 0 i,String.sub s (i+1) (String.length s - i - 1)
        with
          Invalid_argument _ -> 
            c_v := !c_v + 1;
            String.sub s 0 i,Printf.sprintf "unknown_value_%i" !c_k
      with
        | Not_found -> 
           c_k := !c_k + 1;
           (Printf.sprintf "unkown_key_%i" !c_k,s)
    in
    vendor,Array.to_list (Array.map split x)

end

module Encoder = 
struct

  type t

  external init : internal_mode -> int -> t = "ocaml_speex_enc_init" 

  let init m x = init (internal_mode_of_mode m) x

  external get : t -> int -> int = "ocaml_speex_encoder_ctl_get"
  
  let get x q = 
    get x (int_of_control q)

  external set : t -> int -> int -> unit = "ocaml_speex_encoder_ctl_set"

  let set x q v = 
    set x (int_of_control q) v

  external encode_page_main : t -> int -> Ogg.Stream.stream -> (unit -> float array) -> Ogg.Page.t = "ocaml_speex_encode_page"

  let encode_page e s f = encode_page_main e 1 s f

  let merge frame = 
    let ret = Array.init
      (2*(Array.length (frame.(0))))
      (fun _ -> frame.(0).(0))
    in
    Array.iteri
      (fun i ->
         fun _ ->
           ret.(2*i) <- frame.(0).(i);
           ret.(2*i+1) <- frame.(1).(i)
      )  frame.(0);
    ret

  let encode_page_stereo e s f = 
    let f () = merge (f ()) in
     encode_page_main e 2 s f 

  external encode_page_int_main : t -> int -> Ogg.Stream.stream -> (unit -> int array) -> Ogg.Page.t = "ocaml_speex_encode_page_int"

  let encode_page_int e s f = encode_page_int_main e 1 s f

  let encode_page_int_stereo e s f = 
    let f () = merge (f ()) in
    encode_page_int_main e 2 s f 

  external eos : t -> Ogg.Stream.stream -> unit = "ocaml_speex_encoder_eos" 

end

module Decoder = 
struct

  type t
  
  external init : internal_mode -> t = "ocaml_speex_dec_init"

  let init m = init (internal_mode_of_mode m)

  external get : t -> int -> int = "ocaml_speex_decoder_ctl_get"

  let get x q =
    get x (int_of_control q)

  external set : t -> int -> int -> unit = "ocaml_speex_decoder_ctl_set"

  let set x q v =
    set x (int_of_control q) v

  let split frame = 
    let ret = Array.make 2 
       (Array.make (Array.length frame / 2) frame.(0))
    in
    for i = 0 to (Array.length frame / 2) - 1 do
      ret.(0).(i) <- frame.(2*i);
      ret.(1).(i) <- frame.(2*i+1)
    done;
    ret

  external decode_feed : t -> int -> Ogg.Stream.stream -> (float array -> unit) -> unit = "ocaml_speex_decoder_decode"

  let decode_gen e s chan func split = 
    let l = ref [] in
    let feed x = 
      l := split x :: !l
    in
    begin
      try
        func e chan s feed
      with
        | Ogg.Not_enough_data -> 
            if List.length !l = 0 then
              raise Ogg.Not_enough_data
    end;
    List.rev !l

  let decode e s = decode_gen e s 1 decode_feed (fun x -> x)

  let decode_stereo e s = decode_gen e s 2 decode_feed split

  let decode_feed_stereo e s feed =
    let feed x = feed (split x) in 
    decode_feed e 2 s feed

  let decode_feed e s feed =
    decode_feed e 1 s feed

  external decode_int_feed : t -> int -> Ogg.Stream.stream -> (int array -> unit) -> unit = "ocaml_speex_decoder_decode_int"

  let decode_int e s = decode_gen e s 2 decode_int_feed (fun x -> x)

  let decode_int_stereo e s = decode_gen e s 2 decode_int_feed split

  let decode_int_feed_stereo e s feed =
    let feed x = feed (split x) in
    decode_int_feed e 2 s feed

  let decode_int_feed e s feed =
    decode_int_feed e 1 s feed

end

module Wrapper = 
struct

  module Decoder = 
  struct
    
    exception Not_speex
    exception Internal

    type t = (Decoder.t*Ogg.Stream.stream*Ogg.Sync.t*nativeint*int*(string*(string*string) list)*Header.t) ref

    let open_sync sync = 
       (** Test wether the stream contains speex data *)
       let test_speex () =
         (** Get First page *)
         let page = Ogg.Sync.read sync in
         (** Check wether this is a b_o_s *)
         if not (Ogg.Page.bos page) then raise Not_found ;
         (** Create a stream with this ID *)
         let serial = Ogg.Page.serialno page in
         let os = Ogg.Stream.create ~serial () in
         Ogg.Stream.put_page os page ;
         (** Test header. Do not catch anything, first page should be sufficient *)
         let packet = Ogg.Stream.get_packet os in
         try
           let header = Header.header_of_packet packet in
           (* Get comments *)
           let page = ref (Ogg.Sync.read sync) in
           while Ogg.Page.serialno !page <> serial do
             page := Ogg.Sync.read sync
           done;
           Ogg.Stream.put_page os !page ;
           let comments = Ogg.Stream.get_packet os in
           let comments = Header.comments_of_packet comments in
          serial,os,header,comments
         with
           | _ -> Printf.printf "Not a speex stream..\n"; raise Internal
       in
       let rec init () =
         try
          test_speex ()
         with
          (** Not_found is not catched: ogg stream always start
              with all b_o_s and we don't care about sequenced streams here *)
           | Internal -> init () 
           | Ogg.Not_enough_data -> raise Not_speex
       in
       let serial,os,header,comments = init () in
       let chans = header.Header.nb_channels in
       let rate = header.Header.rate in
       let mode = header.Header.mode in
       let dec = Decoder.init mode in
       Decoder.set dec SPEEX_SET_SAMPLING_RATE rate;
       ref (dec,os,sync,serial,chans,comments,header)

    let open_file file = 
      let sync,fd = Ogg.Sync.create_from_file file in
      open_sync sync,fd

    let open_feed feed = 
      let sync = Ogg.Sync.create feed in
      open_sync sync 

    let serial x = 
      let (_,_,_,serial,_,_,_) = !x in
      serial

    let comments x = 
      let (_,_,_,_,_,(_,l),_) = !x in
      l

    let header x = 
      let (_,_,_,_,_,_,h) = !x in
      h

    let decode_gen f x =
    let (dec,os,sync,serial,_,_,_) = !x in 
    let dec = ref dec in
    let serial = ref serial in
    let os = ref os in
    let eos = ref false in
    let rec put () =
      try
        if !eos then
           (* Try to open a new stream *)
         ( try
             let y = open_sync sync in
             let (ndec,nos,_,nserial,_,_,_) = !y in
             x := !y;
             dec := ndec;
             serial := nserial;
             os := nos;
             eos := false
          with
            | Not_found -> raise Internal );
        let page = Ogg.Sync.read sync in
        if Ogg.Stream.eos !os then eos := true;
        if Ogg.Page.serialno page = !serial then
        Ogg.Stream.put_page !os page
      with
        | Internal
        | Ogg.Not_enough_data -> raise Ogg.End_of_stream
        | Ogg.End_of_stream -> eos := true; put ()
    in
    let rec get () = 
      try
        f !dec !os
      with
        | Ogg.End_of_stream -> eos := true; put (); get ()
        | Ogg.Not_enough_data -> put (); get ()
    in
    get ()

    let decode v = decode_gen Decoder.decode v

    let decode_stereo v = decode_gen Decoder.decode_stereo v

    let decode_feed v feed = decode_gen (fun x -> fun y -> Decoder.decode_feed x y feed) v

    let decode_feed_stereo v feed = decode_gen (fun x -> fun y -> Decoder.decode_feed_stereo x y feed) v

    let decode_int v = decode_gen Decoder.decode_int v

    let decode_int_stereo v = decode_gen Decoder.decode_int_stereo v

    let decode_int_feed v feed = decode_gen (fun x -> fun y -> Decoder.decode_int_feed x y feed) v

    let decode_int_feed_stereo v feed = decode_gen (fun x -> fun y -> Decoder.decode_int_feed_stereo x y feed) v


  end
end

module Skeleton = 
struct

  external fisbone : nativeint -> Header.t -> Int64.t -> string -> Ogg.Stream.packet = "ocaml_speex_skeleton_fisbone"

  let fisbone ?(start_granule=Int64.zero)
              ?(headers=["Content-type","audio/speex"])
              ~serialno ~header () =
    let concat s (h,v) =
      Printf.sprintf "%s%s: %s\r\n" s h v
    in
    let s =
      List.fold_left concat "" headers
    in
    fisbone serialno header start_granule s  

end
