exception Error of string

open Caml_cduce
open Asttypes
open Types

(* Unfolding of OCaml types *)

exception PolyAbstract of string

let ocaml_env = ref Env.initial

type t = { uid : int; mutable recurs : int; mutable def : def }
and def =
  | Link of t
  | Arrow of string * t * t
  | Tuple of t list
  | PVariant of (string * t option) list  (* Polymorphic variant *)
  | Variant of string * (string * t list) list * bool
  | Record of string * (string * t) list * bool
  | Builtin of string * t list
  | Abstract of string
  | Var of int

module IntMap = 
  Map.Make(struct type t = int let compare : t -> t -> int = compare end)
module StringMap = 
  Map.Make(struct type t = string let compare : t -> t -> int = compare end)

let rec print_sep f sep ppf = function
  | [] -> ()
  | [x] -> f ppf x
  | x::tl -> Format.fprintf ppf "%a%s" f x sep; print_sep f sep ppf tl

let printed = ref IntMap.empty

let rec print_slot ppf slot =
  if slot.recurs > 0 then
    (
      if IntMap.mem slot.uid !printed then
	Format.fprintf ppf "X%i" slot.uid
      else (
	printed := IntMap.add slot.uid () !printed;
	Format.fprintf ppf "X%i:=%a" slot.uid print_def slot.def
      )
    )
  else 
    print_def ppf slot.def

and print_def ppf = function
  | Link t -> print_slot ppf t
  | Arrow (l,t,s) -> Format.fprintf ppf "%s:%a -> %a" l print_slot t print_slot s
  | Tuple tl -> Format.fprintf ppf "(%a)" (print_sep print_slot ",") tl
  | PVariant l -> Format.fprintf ppf "[%a]" (print_sep print_palt " | ") l
  | Variant (p,l,_) -> Format.fprintf ppf "[%s:%a]" p (print_sep print_alt " | ") l
  | Record (p,l,_) -> Format.fprintf ppf "{%s:%a}" p (print_sep print_field " ; ") l
  | Builtin (p,tl) -> Format.fprintf ppf "%s(%a)" p (print_sep print_slot ",") tl
  | Abstract s -> Format.fprintf ppf "%s" s
  | Var i -> Format.fprintf ppf "'a%i" i


and print_palt ppf = function
  | lab, None -> Format.fprintf ppf "`%s" lab
  | lab, Some t -> Format.fprintf ppf "`%s of %a" lab print_slot t
   
and print_alt ppf = function
  | (lab,[]) ->
      Format.fprintf ppf "%s" lab
  | (lab,l) ->
      Format.fprintf ppf "%s of [%a]" lab (print_sep print_slot ",") l

and print_field ppf (lab,t) =
  Format.fprintf ppf "%s:%a" lab print_slot t


let print = print_slot

let counter = ref 0
let new_slot () =
  incr counter;
  { uid = !counter; recurs = 0; def = Abstract "DUMMY" }

let builtins =
  List.fold_left (fun m x -> StringMap.add x () m) StringMap.empty
    [
      "list"; "Pervasives.ref"; 
      "unit"; "array";
      "Big_int.big_int";
      "Cduce_lib.Value.t"; "Cduce_lib.Encodings.Utf8.t"
    ]

let vars = ref []

let get_var id =
  try List.assq id !vars
  with Not_found ->
    let i = List.length !vars in
    vars := (id,i) :: !vars;
    i

let rec unfold seen constrs ty =
  try
    let t = IntMap.find ty.id seen in
    t.recurs <- t.recurs + 1;
    t
  with Not_found ->
    let slot = new_slot () in
    let seen = IntMap.add ty.id slot seen in
    let loop = unfold seen constrs in
    slot.def <-
    (match ty.desc with
       | Tarrow (l,t1,t2,_) -> Arrow (l, loop t1, loop t2)
       | Ttuple tyl -> Tuple (List.map loop tyl)
       | Tvariant rd ->
	   let fields = 
	     List.fold_left
	       (fun accu (lab,f) -> 
		  match f with
		    | Rpresent (Some t) 
		    | Reither(true, [t], _, _) -> (lab, Some (loop t)) :: accu
		    | Rpresent None 
		    | Reither(true, [], _, _) -> (lab, None) :: accu
		    | Rabsent -> Printf.eprintf "Warning: Rabsent not supported"; accu
		    | Reither _ -> Printf.eprintf "Warning: Reither not supported"; accu
	       ) []
	       rd.row_fields in
	   PVariant fields
       | Tvar -> Var (get_var ty.id)
       | Tconstr (p,args,_) ->
	   let args = List.map loop args in
	   let pn = Path.name p in
	   if StringMap.mem pn builtins 
	   then Builtin (pn,args)
	   else
	     let decl = 
	       try Env.find_type p !ocaml_env
	       with Not_found ->
		 failwith ("Cannot resolve path " ^ pn) in
	     (try
		let (s,args') = StringMap.find pn constrs in
		List.iter2 
		  (fun a a' -> 
		     if a.uid != a'.uid then 
		       failwith "Polymorphic recursion forbidden") args args';
		s.recurs <- s.recurs + 1;
		Link s
	      with Not_found ->
		let seen = 
		  List.fold_left2 
		    (fun seen a v -> a.recurs <- a.recurs - 1; IntMap.add v.id a seen)
		    seen args decl.type_params in
		let constrs = StringMap.add pn (slot,args) constrs in
		let loop = unfold seen constrs in
		let prefix = match p with
		  | Path.Pident _ -> ""
		  | Path.Pdot (p,_,_) -> Path.name p ^ "."
		  | _ -> assert false in
		(match decl.type_kind, decl.type_manifest with
		   | Type_variant (cstrs,pub), _ ->
		       let cstrs =
			 List.map (fun (cst,f) -> (cst,List.map loop f)) cstrs in
		       Variant (prefix, cstrs, pub = Public)
		   | Type_record (f,_,pub), _ ->
		       let f = List.map (fun (l,_,t) -> (l,loop t)) f in
		       Record (prefix, f, pub = Public)
		   | Type_abstract, Some t ->
		       Link (loop t)
		   | Type_abstract, None ->
		       (match args with
			  | [] -> Abstract pn
			  | _ -> raise (PolyAbstract pn))))
       | _ -> failwith "Unsupported feature"
    );
    slot

let unfold ty = 
  vars := [];
  let t = unfold IntMap.empty StringMap.empty ty in
  let n = List.length !vars in
  vars := [];
  (t,n)

(* Reading .cmi *)

let unsupported s =
  raise (Error (Printf.sprintf "Unsupport feature (%s) found in .cmi" s))

let has_cmi name =
  Config.load_path := Config.standard_library :: !Librarian.obj_path;
  try ignore (Misc.find_in_path_uncap !Config.load_path (name ^ ".cmi")); true
  with Not_found -> false

let load_cmi name = 
  Config.load_path := Config.standard_library :: !Librarian.obj_path;
  let filename = Misc.find_in_path_uncap !Config.load_path (name ^ ".cmi") in
  let sg = Env.read_signature name filename in
  ocaml_env := Env.add_signature sg Env.initial;
  let values = ref [] in
  List.iter
    (function
       | Tsig_value (id, {val_type=t;val_kind=Val_reg|Val_prim _}) -> 
	   let n = name ^ "." ^ (Ident.name id) in
	   (try
	     let (t,_) = unfold t in
	     values := (name ^ "." ^ (Ident.name id), t) :: !values
	   with PolyAbstract _ -> ())
       | _ -> ()
    ) sg;
  !values

let load_cmi name =
  try load_cmi name
  with Env.Error e ->
    Env.report_error Format.str_formatter e;
    let s = Format.flush_str_formatter () in
    let s = Printf.sprintf "Error while reading OCaml interface %s: %s"
	      name s in
    raise (Location.Generic s)

let read_cmi name =
  Config.load_path := Config.standard_library :: !Librarian.obj_path;
  let filename = Misc.find_in_path_uncap !Config.load_path (name ^ ".cmi") in
  let sg = Env.read_signature name filename in
  ocaml_env := Env.add_signature sg Env.initial;
  let buf = Buffer.create 1024 in
  let ppf = Format.formatter_of_buffer buf in
  let values = ref [] in
  List.iter
    (function
       | Tsig_value (id, {val_type=t;val_kind=Val_reg}) -> 
	   let (unf,n) = unfold t in
	   if n !=0 then unsupported "polymorphic value";
	   values := (Ident.name id, t, unf) :: !values
       | Tsig_type (id,t,_) -> 
	   Format.fprintf ppf "%a@." (Printtyp.type_declaration id) t
       | Tsig_value _ -> unsupported "external value"
       | Tsig_exception _ -> unsupported "exception"
       | Tsig_module _ -> unsupported "module"
       | Tsig_modtype _ -> unsupported "module type"
       | Tsig_class _ -> unsupported "class"
       | Tsig_cltype _ -> unsupported "class type"
    ) sg;
  (Buffer.contents buf, !values)

let read_cmi name =
  try read_cmi name
  with Env.Error e ->
    Env.report_error Format.str_formatter e;
    let s = Format.flush_str_formatter () in
    let s = Printf.sprintf "Error while reading OCaml interface %s: %s"
	      name s in
    raise (Location.Generic s)


let print_ocaml = Printtyp.type_expr


let rec dump_li = function
  | Longident.Lident s -> print_endline s
  | Longident.Ldot (li,s) -> dump_li li; print_endline s
  | _ -> assert false

let find_value v =
  Config.load_path := Config.standard_library :: !Librarian.obj_path;
  let li = Longident.parse v in
  let (p,vd) = Env.lookup_value li Env.initial in
  unfold vd.val_type