(* based on code pasted by pango on FreeNode/#ocaml *)

open Printf
open Scanf

let debug = false

type annot_t = {
  annot_begin : int;
  annot_end : int;
  annot_type : string
}

let parse_annotations annot =
  let sbuf = Scanning.from_file annot in
  let result = ref [] in
  try
    while true do
      let boffset, eoffset =
	bscanf sbuf " %S %d %d %d %S %d %d %d" (fun _ _ _ boffset _ _ _ eoffset -> (boffset, eoffset)) in
      ignore(bscanf sbuf " type(" (fun () -> ()) ());
      let rec read_signature acc =
	let line = Scanf.bscanf sbuf " %s@\n" (fun s -> s) in
	if line = ")" then List.rev acc
	else read_signature (line :: acc) in
      let signature = String.concat " " (read_signature []) in
      if debug then eprintf "%d-%d: %s\n" boffset eoffset signature;
      result := { annot_begin = boffset; annot_end = eoffset; annot_type = signature } :: !result
    done;
    assert false
  with
  | End_of_file
  | Scan_failure _ -> !result


let merge source annot dest =
  let annotations = parse_annotations annot in
  let beginnings = Hashtbl.create 10 in
  let endings = Hashtbl.create 10 in
  List.iter (fun a ->
	       Hashtbl.add beginnings a.annot_begin a;
	       Hashtbl.add endings a.annot_end a
	    ) annotations;
  let source = open_in source in
  try
    let rec loop i =
      List.iter (fun a -> dest "(") (Hashtbl.find_all beginnings i);
      dest (Printf.sprintf "%c" (input_char source));
      List.iter (fun a -> dest (Printf.sprintf ":%s)" a.annot_type))
	(List.sort (fun a b -> compare b.annot_begin a.annot_begin)
	   (Hashtbl.find_all endings (i+1)));
      loop (i+1) in
    loop 0
  with End_of_file ->
    close_in source

let annotator_main () =
  let source = ref (None : string option) in
  let annot = ref (None : string option) in
  Arg.parse [
    ("-source", Arg.String (fun s -> source := Some s), "name of module to annotate");
    ("-annot", Arg.String (fun a -> annot := Some a), "name of the annotation file to merge")
  ]
    (fun s -> eprintf "Don't know what to do with %s, ignoring\n%!" s)
    "Merge annotations (produced by ocamlc/ocamlopt -dtypes) with module source";
  
  match !source, !annot with
  | Some s, Some a -> merge s a print_string
  | _ ->
      eprintf "Arguments missing\n";
      exit 1