open Types

exception OggUnalignedStream;;

let identify_bos page = 
    let fst_pkt = List.hd page.raw_data in
    if String.sub fst_pkt 0 7 = "fishead" then Skeleton
    else if String.sub fst_pkt 0 4 = "CMML" then CMML
    else if String.sub fst_pkt 1 6 = "theora" then Theora
    else if String.sub fst_pkt 1 6 = "vorbis" then Vorbis
    else Unknown

(*
 * generate a packet size from the stream (n lacing values left).  INTERNAL
 * FUNCTION
 *)
let rec generate_packet stream n =
  match n with
  | 0 ->  ((0, 1), 0)
  | _ ->  let v = IO.read_byte stream in 
          if v = 255 then
            let ((r, c), u) = generate_packet stream (n - 1) in
            ((r + 255, c), u + 1)
          else
            ((v, 0), 1)

(*
 * generate the packet sizes from the n lacing values.  INTERNAL FUNCTION
 *)
let rec generate_packets stream n =
  match n with
  | 0 -> []
  | _ -> let (v, r) = generate_packet stream n in 
                      v::(generate_packets stream (n - r))

(*
 * extract the packet sizes from the stream.  INTERNAL FUNCTION
 *)
let get_packet_sizes stream = 
  let num_segments = IO.read_byte stream in 
  generate_packets stream num_segments;; 

(*
 * extract packets as a string list from stream, given sizes.  Actually returns
 * (size list, raw packet list, last packet continued).  INTERNAL FUNCTION
 *)
let rec extract_packets pkts stream =
  match pkts with
    | [] -> ([], [], 0)
    | [(s,c)] -> ([s], [IO.really_nread stream s], c)
    | (s,c)::tl -> begin
                    let p = IO.really_nread stream s in
                    let (os, op, c) = extract_packets tl stream in
                      (s::os, p::op, c)
                  end;;

(* 
 * make a page.  INTERNAL FUNCTION
 *)
let make_page h gp sn seq pkts stream cksum =
  let (sizes, packets, lpc) = extract_packets pkts stream in
  { continued = (h land 0x1) > 0;
    bos = (h land 0x2) > 0;
    eos = (h land 0x4) > 0;
    last_packet_complete = (lpc = 0);
    granulepos = gp;
    serialno = sn;
    sequenceno = seq;
    checksum = cksum;
    packet_sizes = sizes;
    raw_data = packets
  };;

let make_header page = 
        if page.continued then 1 else 0 +
        if page.bos then 2 else 0 +
        if page.eos then 4 else 0 
    
let make_checksum page = page.checksum;;

let rec write_lacing_value stream size tpc =
  if (not tpc) && size == 0 then ()
  else if size < 255 then IO.write_byte stream size
  else 
    begin 
      IO.write_byte stream 255; 
      write_lacing_value stream (size - 255) tpc 
    end;;

let rec num_lacing_values sizes lpc = 
  match sizes with
    | [] -> 0
    | [h] -> ((h - if lpc then 0 else 1) / 255) + 1
    | h::t -> (h / 255) + 1 + (num_lacing_values t lpc)

let rec write_lacing_values stream sizes lpc = 
  match sizes with
    | []    ->  ()
    | [h]   ->  write_lacing_value stream h lpc
    | h::t  ->  begin 
                  write_lacing_value stream h true; 
                  write_lacing_values stream t lpc
                end;;
    
let rec write_packet_data stream data = 
  match data with
    | []    ->  ()
    | h::t  ->  begin IO.nwrite stream h; write_packet_data stream t end;;

let packets_to_stream stream page = 
  let sizes = page.packet_sizes in
  let data = page.raw_data in
  let lpc = page.last_packet_complete in 
  IO.write_byte stream (num_lacing_values sizes lpc);
  write_lacing_values stream sizes lpc;
  write_packet_data stream data;;

let wrap_value = Int64.sub (Int64.shift_left Int64.one 32) Int64.one;;
 
(*
let read_ui32 stream =
  let b = IO.read_ui16 stream in
  let a = IO.read_ui16 stream in
  Int64.add (Int64.shift_left (Int64.of_int a) 16) (Int64.of_int b);;

let mask = Int64.of_int 0xFFFF
  
let write_ui32 stream value =
  let a = Int64.to_int (Int64.shift_right value 16) in
  let b = Int64.to_int (Int64.logand value mask) in
  IO.write_ui16 stream b;
  IO.write_ui16 stream a;;

let read_oogg64 stream = 
    (IO.read_ui16 stream, IO.read_ui16 stream, IO.read_ui16 stream,
        IO.read_ui16 stream);;

let write_oogg64 stream (a, b, c, d) = 
    IO.write_ui16 stream d; IO.write_ui16 stream c; 
    IO.write_ui16 stream b; IO.write_ui16 stream a;;

let read_oogg32 stream = 
    (IO.read_ui16 stream, IO.read_ui16 stream);;

let write_oogg32 stream (a, b) =
    IO.write_ui16 stream b; IO.write_ui16 stream a;;

let read_granulepos stream = 
  let gp = read_oogg64 stream in 
  if gp = (0xFFFF,0xFFFF,0xFFFF,0xFFFF) then None else Some gp;;

let write_granulepos stream gp = match gp with 
  | None -> write_oogg64 stream (0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF)
  | Some gp -> write_oogg64 stream gp;;

*)

let page_to_raw stream page = 
  IO.nwrite stream "OggS"; IO.write_byte stream 0;
  IO.write_byte stream (make_header page);
  write_granulepos stream page.granulepos;
  write_oogg32 stream page.serialno;
  write_oogg32 stream page.sequenceno;
  write_oogg32 stream (make_checksum page);
  packets_to_stream stream page;;
  
(*
 * generate a page stream from a raw stream
 *)
let rec to_rawPageStream stream = 
  try begin
    let header = IO.nread stream 4 in
    if header = "OggS" then
      begin
        ignore (IO.read_byte stream);
        let header_type = IO.read_byte stream in
        let granule_pos = read_granulepos stream in
        let serial_no = read_oogg32 stream in
        let seq_no = read_oogg32 stream in
        let checksum = read_oogg32 stream in
        let packets = get_packet_sizes stream in
        [< 'make_page header_type granule_pos serial_no seq_no packets stream 
              checksum; to_rawPageStream stream >]
      end
    else
      raise OggUnalignedStream
  end with IO.No_more_input -> [< >];;

(*
 * header: 4, granulepos: 8, serialno: 4, seqno: 4, checksum: 4, 
 * num_lacing_val: 1, lacing_vals: num_lacing_values, packets: sum(sizes)
 *
 * total: 25 + num_lacing_values + sum(sizes)
 *)

let oogg_page_to_string page =
  let the_string = IO.output_string () in page_to_raw the_string page;
  IO.close_out the_string;;

type contextRecord = { mutable ct_info : 
      (serialNo * ((granulePos -> float option) * mediaType)) list };;

let add_bos context rpage =
  let sn = rpage.serialno in
  let id = identify_bos rpage in
  let bospack = List.hd rpage.raw_data in
  context.ct_info <- 
          (sn, (Granules.granulerate_function id bospack, id))::context.ct_info;
  id;;

let lookup_bos context rpage =
  let (f, id) = List.assoc rpage.serialno context.ct_info in
  (id, f rpage.granulepos)

let rawPage_to_page context rpage = 
  if rpage.bos then (
    let id = add_bos context rpage in
    { raw = rpage; time = Some 0.0; identity = id }
  ) else (
    let (id, time) = lookup_bos context rpage in
    { raw = rpage; time = time; identity = id }
  );;

(*
 * write a page stream to a raw stream
 *)
let rec write_rawPageStream stream pstream = 
  match pstream with parser
   | [< 'page ; rest >] -> begin 
                             page_to_raw stream page; 
                             write_rawPageStream stream rest
                           end
   | [< >] -> ();;

let rec check_crc stream =
  match stream with parser
    | [< 'page ; rest >] -> 
      begin
        let crc = page.checksum in
        let newpage = {page with checksum=(0, 0)} in
        let compcrc = Crc.crc (oogg_page_to_string newpage) in
        if not (compcrc = crc) then begin
          print_oogg32 crc ; print_string " " ; 
          print_oogg32 compcrc ; print_newline () end;
        check_crc rest
      end
    | [< >] -> ();;

let generate_crc page = Crc.crc (oogg_page_to_string page);;

let rawPageStream_to_pageStream rpstream = 
  let context = { ct_info = [] } in
  let rec _rps_to_ps rpstream = match rpstream with parser
    | [< 'rpage ; rest >] -> 
          [< 'rawPage_to_page context rpage; _rps_to_ps rest >]
    | [< >] -> [< >] in
  _rps_to_ps rpstream;;

let to_pageStream i = rawPageStream_to_pageStream (to_rawPageStream i);;

let rec pageStream_to_rawPageStream pstream = match pstream with parser
  | [< 'page ; rest >] -> [< 'page.raw; pageStream_to_rawPageStream rest >]
  | [< >] -> [< >];;

let write_pageStream stream pstream = 
  write_rawPageStream stream (pageStream_to_rawPageStream pstream);;

module RawPageInput =
  struct
    exception DontAttemptToMergeRawStreams
    type k = serialNo
    type s = rawPage
    let eq = (=)
    let get_key p = p.serialno
    let is_first p = p.bos
    let is_last p = p.eos
    let printKey s = print_oogg32 s
    let lt a b = raise DontAttemptToMergeRawStreams
  end;;

module RawPageSort = StreamSort.StreamSort (RawPageInput);;
let to_streams_raw = RawPageSort.sort;;

module PageInput =
  struct
    type k = serialNo
    type s = page
    let eq = (=)
    let get_key p = p.raw.serialno
    let is_first p = p.raw.bos
    let is_last p = p.raw.eos
    let printKey s = print_oogg32 s
    let lt a b = 
      if a.raw.bos then true
      else if b.raw.bos then false
      else match a.time with
      | None -> true
      | Some t1 -> (
        match b.time with
          | None -> false
          | Some t2 -> t1 < t2)
  end;;

module PageSort = StreamSort.StreamSort (PageInput);;
let to_streams = PageSort.sort;;
let from_streams = PageSort.merge;;

let _cmp a b = 
  if a = None then true
  else if b = None then true
  else a >= b;;

let _p t = match t with
  | None -> "None"
  | Some t -> "Some " ^ string_of_float t

let sort pageStream =
  let sorted = ref false in
  let lookaside = ref None in
  let last = ref (Some 0.0) in
  let rec _sort pageStream = match pageStream with parser
    | [< 'page >] -> 
      (
        if !last > page.time then sorted := false;
        match !lookaside with
          | None -> lookaside := Some page; [< _sort pageStream >]
          | Some stored -> 
            (
              if _cmp page.time stored.time
              then 
                (
                  lookaside := Some page; 
                  last := stored.time ; 
                  [< 'stored ; _sort pageStream >]
                )
              else 
                (
                  last := page.time ; 
                  [< 'page ; _sort pageStream >]
                )
            )
      )
    | [< >] -> 
        ( match !lookaside with
          | None -> [< >]
          | Some stored -> [< 'stored >] ) in
  let ps = ref pageStream in
  while not !sorted do (
    sorted := true;
    ps := _sort !ps;
  ) done;
  !ps;;