module Error_monad = X_error_monad
open Tezos_micheline

let env = Error_monad.force_lwt ~msg:"Cast:init environment" @@ Init_proto_alpha.init_environment ()

open Memory_proto_alpha
open Protocol
open Alpha_context

exception Expr_from_string
let expr_of_string str =
  let (ast, errs) = Michelson_parser.V1.parse_expression ~check:false str in
  (match errs with
   | [] -> ()
   | lst -> (
       Format.printf "expr_from_string: %a\n" Error_monad.pp_print_error lst;
       raise Expr_from_string
     ));
  ast.expanded

let tl_of_string str =
  let (ast, errs) = Michelson_parser.V1.parse_toplevel ~check:false str in
  (match errs with
   | [] -> ()
   | lst -> (
       Format.printf "expr_from_string: %a\n" Error_monad.pp_print_error lst;
       raise Expr_from_string
     ));
  ast.expanded

let lexpr_of_string str =
  Script.lazy_expr @@ expr_of_string str

let ltl_of_string str =
  Script.lazy_expr @@ tl_of_string str

let node_of_string str =
  Micheline.root @@ expr_of_string str

let node_to_string (node:_ Micheline.node) =
  let stripped = Micheline.strip_locations node in
  let print_node = Micheline_printer.printable Michelson_v1_primitives.string_of_prim stripped in
  Micheline_printer.print_expr Format.str_formatter print_node ;
  Format.flush_str_formatter ()

open Script_ir_translator

type ex_typed_value =
  Ex_typed_value : ('a Script_typed_ir.ty * 'a) -> ex_typed_value

include struct
  open Script_typed_ir
  open Protocol.Environment.Error_monad
  module Unparse_costs = Michelson_v1_gas.Cost_of.Unparse
  open Micheline
  open Michelson_v1_primitives
  open Protocol.Environment



  let rec unparse_data_generic
    : type a. context -> ?mapper:_ -> unparsing_mode -> a ty -> a -> (Script.node * context) tzresult Lwt.t
    = fun ctxt ?(mapper = fun _ -> return None) mode ty a ->
      Lwt.return (Gas.consume ctxt Unparse_costs.cycle) >>=? fun ctxt ->
      mapper (Ex_typed_value (ty, a)) >>=? function
      | Some x -> return (x , ctxt)
      | None -> (
          match ty, a with
          | Unit_t _, () ->
            Lwt.return (Gas.consume ctxt Unparse_costs.unit) >>=? fun ctxt ->
            return (Prim (-1, D_Unit, [], []), ctxt)
          | Int_t _, v ->
            Lwt.return (Gas.consume ctxt (Unparse_costs.int v)) >>=? fun ctxt ->
            return (Int (-1, Script_int.to_zint v), ctxt)
          | Nat_t _, v ->
            Lwt.return (Gas.consume ctxt (Unparse_costs.int v)) >>=? fun ctxt ->
            return (Int (-1, Script_int.to_zint v), ctxt)
          | String_t _, s ->
            Lwt.return (Gas.consume ctxt (Unparse_costs.string s)) >>=? fun ctxt ->
            return (String (-1, s), ctxt)
          | Bytes_t _, s ->
            Lwt.return (Gas.consume ctxt (Unparse_costs.bytes s)) >>=? fun ctxt ->
            return (Bytes (-1, s), ctxt)
          | Bool_t _, true ->
            Lwt.return (Gas.consume ctxt Unparse_costs.bool) >>=? fun ctxt ->
            return (Prim (-1, D_True, [], []), ctxt)
          | Bool_t _, false ->
            Lwt.return (Gas.consume ctxt Unparse_costs.bool) >>=? fun ctxt ->
            return (Prim (-1, D_False, [], []), ctxt)
          | Timestamp_t _, t ->
            Lwt.return (Gas.consume ctxt (Unparse_costs.timestamp t)) >>=? fun ctxt ->
            begin
              match mode with
              | Optimized -> return (Int (-1, Script_timestamp.to_zint t), ctxt)
              | Readable ->
                match Script_timestamp.to_notation t with
                | None -> return (Int (-1, Script_timestamp.to_zint t), ctxt)
                | Some s -> return (String (-1, s), ctxt)
            end
          | Address_t _, (c, entrypoint)  ->
            Lwt.return (Gas.consume ctxt Unparse_costs.contract) >>=? fun ctxt ->
            begin
              match mode with
              | Optimized ->
                let entrypoint = match entrypoint with "default" -> "" | name -> name in
                let bytes = Data_encoding.Binary.to_bytes_exn
                    Data_encoding.(tup2 Contract.encoding Variable.string)
                    (c, entrypoint) in
                return (Bytes (-1, bytes), ctxt)
              | Readable ->
                let notation = match entrypoint with
                  | "default" -> Contract.to_b58check c
                  | entrypoint -> Contract.to_b58check c ^ "%" ^ entrypoint in
                return (String (-1, notation), ctxt)
            end
          | Contract_t _, (_, (c, entrypoint))  ->
            Lwt.return (Gas.consume ctxt Unparse_costs.contract) >>=? fun ctxt ->
            begin
              match mode with
              | Optimized ->
                let entrypoint = match entrypoint with "default" -> "" | name -> name in
                let bytes = Data_encoding.Binary.to_bytes_exn
                    Data_encoding.(tup2 Contract.encoding Variable.string)
                    (c, entrypoint) in
                return (Bytes (-1, bytes), ctxt)
              | Readable ->
                let notation = match entrypoint with
                  | "default" -> Contract.to_b58check c
                  | entrypoint -> Contract.to_b58check c ^ "%" ^ entrypoint in
                return (String (-1, notation), ctxt)
            end
          | Signature_t _, s ->
            Lwt.return (Gas.consume ctxt Unparse_costs.signature) >>=? fun ctxt ->
            begin
              match mode with
              | Optimized ->
                let bytes = Data_encoding.Binary.to_bytes_exn Signature.encoding s in
                return (Bytes (-1, bytes), ctxt)
              | Readable ->
                return (String (-1, Signature.to_b58check s), ctxt)
            end
          | Mutez_t _, v ->
            Lwt.return (Gas.consume ctxt Unparse_costs.tez) >>=? fun ctxt ->
            return (Int (-1, Z.of_int64 (Tez.to_mutez v)), ctxt)
          | Key_t _, k ->
            Lwt.return (Gas.consume ctxt Unparse_costs.key) >>=? fun ctxt ->
            begin
              match mode with
              | Optimized ->
                let bytes = Data_encoding.Binary.to_bytes_exn Signature.Public_key.encoding k in
                return (Bytes (-1, bytes), ctxt)
              | Readable ->
                return (String (-1, Signature.Public_key.to_b58check k), ctxt)
            end
          | Key_hash_t _, k ->
            Lwt.return (Gas.consume ctxt Unparse_costs.key_hash) >>=? fun ctxt ->
            begin
              match mode with
              | Optimized ->
                let bytes = Data_encoding.Binary.to_bytes_exn Signature.Public_key_hash.encoding k in
                return (Bytes (-1, bytes), ctxt)
              | Readable ->
                return (String (-1, Signature.Public_key_hash.to_b58check k), ctxt)
            end
          | Operation_t _, (op, _big_map_diff) ->
            let bytes = Data_encoding.Binary.to_bytes_exn Alpha_context.Operation.internal_operation_encoding op in
            Lwt.return (Gas.consume ctxt (Unparse_costs.operation bytes)) >>=? fun ctxt ->
            return (Bytes (-1, bytes), ctxt)
          | Chain_id_t _, chain_id ->
            let bytes = Data_encoding.Binary.to_bytes_exn Chain_id.encoding chain_id in
            Lwt.return (Gas.consume ctxt (Unparse_costs.chain_id bytes)) >>=? fun ctxt ->
            return (Bytes (-1, bytes), ctxt)
          | Pair_t ((tl, _, _), (tr, _, _), _, _), (l, r) ->
            Lwt.return (Gas.consume ctxt Unparse_costs.pair) >>=? fun ctxt ->
            unparse_data_generic ctxt mode tl l >>=? fun (l, ctxt) ->
            unparse_data_generic ctxt mode tr r >>=? fun (r, ctxt) ->
            return (Prim (-1, D_Pair, [ l; r ], []), ctxt)
          | Union_t ((tl, _), _, _, _), L l ->
            Lwt.return (Gas.consume ctxt Unparse_costs.union) >>=? fun ctxt ->
            unparse_data_generic ctxt mode tl l >>=? fun (l, ctxt) ->
            return (Prim (-1, D_Left, [ l ], []), ctxt)
          | Union_t (_, (tr, _), _, _), R r ->
            Lwt.return (Gas.consume ctxt Unparse_costs.union) >>=? fun ctxt ->
            unparse_data_generic ctxt mode tr r >>=? fun (r, ctxt) ->
            return (Prim (-1, D_Right, [ r ], []), ctxt)
          | Option_t (t, _, _), Some v ->
            Lwt.return (Gas.consume ctxt Unparse_costs.some) >>=? fun ctxt ->
            unparse_data_generic ctxt mode t v >>=? fun (v, ctxt) ->
            return (Prim (-1, D_Some, [ v ], []), ctxt)
          | Option_t _, None ->
            Lwt.return (Gas.consume ctxt Unparse_costs.none) >>=? fun ctxt ->
            return (Prim (-1, D_None, [], []), ctxt)
          | List_t (t, _, _), items ->
            fold_left_s
              (fun (l, ctxt) element ->
                 Lwt.return (Gas.consume ctxt Unparse_costs.list_element) >>=? fun ctxt ->
                 unparse_data_generic ctxt mode t element >>=? fun (unparsed, ctxt) ->
                 return (unparsed :: l, ctxt))
              ([], ctxt)
              items >>=? fun (items, ctxt) ->
            return (Micheline.Seq (-1, List.rev items), ctxt)
          | Set_t (t, _), set ->
            let t = ty_of_comparable_ty t in
            fold_left_s
              (fun (l, ctxt) item ->
                 Lwt.return (Gas.consume ctxt Unparse_costs.set_element) >>=? fun ctxt ->
                 unparse_data_generic ctxt mode t item >>=? fun (item, ctxt) ->
                 return (item :: l, ctxt))
              ([], ctxt)
              (set_fold (fun e acc -> e :: acc) set []) >>=? fun (items, ctxt) ->
            return (Micheline.Seq (-1, items), ctxt)
          | Map_t (kt, vt, _, _), map ->
            let kt = ty_of_comparable_ty kt in
            fold_left_s
              (fun (l, ctxt) (k, v) ->
                 Lwt.return (Gas.consume ctxt Unparse_costs.map_element) >>=? fun ctxt ->
                 unparse_data_generic ctxt mode kt k >>=? fun (key, ctxt) ->
                 unparse_data_generic ctxt mode vt v >>=? fun (value, ctxt) ->
                 return (Prim (-1, D_Elt, [ key ; value ], []) :: l, ctxt))
              ([], ctxt)
              (map_fold (fun k v acc -> (k, v) :: acc) map []) >>=? fun (items, ctxt) ->
            return (Micheline.Seq (-1, items), ctxt)
          | Big_map_t (kt, vt, _), { id = None ; diff = (module Diff) ; _ } ->
            (* this branch is to allow roundtrip of big map literals *)
            let kt = ty_of_comparable_ty kt in
            fold_left_s
              (fun (l, ctxt) (k, v) ->
                 Lwt.return (Gas.consume ctxt Unparse_costs.map_element) >>=? fun ctxt ->
                 unparse_data_generic ctxt mode kt k >>=? fun (key, ctxt) ->
                 unparse_data_generic ctxt mode vt v >>=? fun (value, ctxt) ->
                 return (Prim (-1, D_Elt, [ key ; value ], []) :: l, ctxt))
              ([], ctxt)
              (Diff.OPS.fold
                 (fun k v acc -> match v with | None -> acc | Some v -> (k, v) :: acc)
                 (fst Diff.boxed) []) >>=? fun (items, ctxt) ->
            return (Micheline.Seq (-1, items), ctxt)
          | Big_map_t (_kt, _kv, _), { id = Some id ; diff = (module Diff) ; _ } ->
            if Compare.Int.(Diff.OPS.cardinal (fst Diff.boxed) = 0) then
              return (Micheline.Int (-1, id), ctxt)
            else
              (* this can only be the result of an execution and the map
                 must have been flushed at this point *)
              assert false
          | Lambda_t _, Lam (_, original_code) ->
            unparse_code_generic ctxt ~mapper mode original_code
        )

  and unparse_code_generic ctxt ?mapper mode =
    let legacy = true in
    function
    | Prim (loc, I_PUSH, [ ty ; data ], annot) ->
      Lwt.return (parse_packable_ty ctxt ~legacy ty) >>=? fun (Ex_ty t, ctxt) ->
      parse_data ctxt ~legacy t data >>=? fun (data, ctxt) ->
      unparse_data_generic ctxt ?mapper mode t data >>=? fun (data, ctxt) ->
      Lwt.return (Gas.consume ctxt (Unparse_costs.prim_cost 2 annot)) >>=? fun ctxt ->
      return (Prim (loc, I_PUSH, [ ty ; data ], annot), ctxt)
    | Seq (loc, items) ->
      fold_left_s
        (fun (l, ctxt) item ->
           unparse_code_generic ctxt ?mapper mode item >>=? fun (item, ctxt) ->
           return (item :: l, ctxt))
        ([], ctxt) items >>=? fun (items, ctxt) ->
      Lwt.return (Gas.consume ctxt (Unparse_costs.seq_cost (List.length items))) >>=? fun ctxt ->
      return (Micheline.Seq (loc, List.rev items), ctxt)
    | Prim (loc, prim, items, annot) ->
      fold_left_s
        (fun (l, ctxt) item ->
           unparse_code_generic ctxt ?mapper mode item >>=? fun (item, ctxt) ->
           return (item :: l, ctxt))
        ([], ctxt) items >>=? fun (items, ctxt) ->
      Lwt.return (Gas.consume ctxt (Unparse_costs.prim_cost 3 annot)) >>=? fun ctxt ->
      return (Prim (loc, prim, List.rev items, annot), ctxt)
    | Int _ | String _ | Bytes _ as atom -> return (atom, ctxt)

end

let rec mapper (Ex_typed_value (ty, a)) =
  let open Alpha_environment.Error_monad in
  let open Script_typed_ir in
  let open Micheline in
  match ty, a with
  | Big_map_t (kt, vt, Some (`Type_annot "toto")), map ->
    let kt = ty_of_comparable_ty kt in
    fold_left_s
      (fun l (k, v) ->
         match v with
         | None -> return l
         | Some v -> (
             let key = data_to_node (Ex_typed_value (kt, k)) in
             let value = data_to_node (Ex_typed_value (vt, v)) in
             return (Prim (-1, Michelson_v1_primitives.D_Elt, [ key ; value ], []) :: l))
      )
      []
      (map_fold (fun k v acc -> (k, v) :: acc) map.diff []) >>=? fun items ->
    return (Some (Micheline.Seq (-1, String (-1, "...") :: items)))
  | _ -> return None

and data_to_node (Ex_typed_value (ty, data)) =
  let tc = env.tezos_context in
  let node_lwt = unparse_data_generic tc ~mapper Readable ty data in
  let node = fst @@ Error_monad.force_lwt_alpha ~msg:"data to string" node_lwt in
  node

let data_to_string ty data =
  let node = data_to_node (Ex_typed_value (ty, data)) in
  node_to_string node

open Script_typed_ir
open Script_interpreter
type ex_typed_stack =
    Ex_typed_stack : ('a stack_ty * 'a stack) -> ex_typed_stack

let stack_to_string stack_ty stack =
  let rec aux acc fst (Ex_typed_stack(stack_ty,stack)) =
    match (stack_ty, stack) with
    | Item_t (hd_ty, tl_ty, _), Item (hd, tl) -> (
        let separator = if not fst then " ; " else "" in
        let str = data_to_string hd_ty hd in
        let acc = acc ^ separator ^ str in
        let new_value = aux acc false (Ex_typed_stack (tl_ty, tl)) in
        new_value
      )
    | _ -> acc in
  aux "" true @@ Ex_typed_stack(stack_ty, stack)

let ty_to_node ty =
  let (node, _) = Error_monad.force_lwt_alpha ~msg:"ty to node" @@ Script_ir_translator.unparse_ty env.tezos_context ty in
  node

type ex_descr =
    Ex_descr : (_, _) Script_typed_ir.descr -> ex_descr

let descr_to_node x =
  let open Alpha_context.Script in
  let open Micheline in
  let open! Script_typed_ir in
  let rec f : ex_descr -> Script.node = fun descr ->
    let prim ?children ?children_nodes p =
      match (children, children_nodes) with
      | Some children, None ->
        Prim (0, p, List.map f children, [])
      | Some _, Some _ ->
        raise @@ Failure "descr_to_node: too many parameters"
      | None, Some children_nodes ->
        Prim (0, p, children_nodes, [])
      | None, None ->
        Prim (0, p, [], [])
    in
    let (Ex_descr descr) = descr in
    match descr.instr with
    | Dup -> prim I_DUP
    | Drop -> prim I_DROP
    | Swap -> prim I_SWAP
    | Dip c -> prim ~children:[Ex_descr c] I_DIP
    | Car -> prim I_CAR
    | Cdr -> prim I_CDR
    | Cons_pair -> prim I_PAIR
    | Nop -> Micheline.Seq (0, [prim I_UNIT ; prim I_DROP])
    | Seq (a, b) -> Micheline.Seq (0, List.map f [Ex_descr a ; Ex_descr b])
    | Const v -> (
        let (Item_t (ty, _, _)) = descr.aft in
        prim  ~children_nodes:[data_to_node (Ex_typed_value (ty, v))] I_PUSH
      )
    | Failwith _ -> prim I_FAILWITH
    | If (a, b) -> prim ~children:[Ex_descr a ; Ex_descr b] I_IF
    | Loop c -> prim ~children:[Ex_descr c] I_LOOP
    | If_left (a, b) -> prim ~children:[Ex_descr a ; Ex_descr b] I_IF_LEFT
    | Left -> prim I_LEFT
    | Right -> prim I_RIGHT
    | Loop_left c -> prim ~children:[Ex_descr c] I_LOOP_LEFT
    | If_none (a, b) -> prim ~children:[Ex_descr a ; Ex_descr b] I_IF_NONE
    | Cons_none _ -> prim I_NONE
    | Cons_some -> prim I_SOME
    | Nil -> prim I_NIL
    | Cons_list -> prim I_CONS
    | If_cons (a, b) -> prim ~children:[Ex_descr a ; Ex_descr b] I_IF_CONS
    | List_iter _ -> prim I_ITER
    | Compare _ -> prim I_COMPARE
    | Int_nat -> prim I_INT
    | Add_natnat -> prim I_ADD
    | Add_natint -> prim I_ADD
    | Add_intnat -> prim I_ADD
    | Sub_int -> prim I_SUB
    | Mul_natnat -> prim I_MUL
    | Ediv_natnat -> prim I_MUL
    | Map_get -> prim I_GET
    | Map_update -> prim I_UPDATE
    | Big_map_get -> prim I_GET
    | Big_map_update -> prim I_UPDATE
    | Gt -> prim I_GT
    | Ge -> prim I_GE
    | Pack _ -> prim I_PACK
    | Unpack _ -> prim I_UNPACK
    | Blake2b -> prim I_BLAKE2B
    | And -> prim I_AND
    | Xor -> prim I_XOR
    | _ -> raise @@ Failure "descr to node" in
  f @@ Ex_descr x

let rec flatten_node =
  let open! Micheline in
  function
  | Micheline.Seq (a, lst) -> (
      let aux = function
        | Prim (loc, p, children, annot) -> [ Prim (loc, p, List.map flatten_node children, annot) ]
        | Seq (_, lst) -> List.map flatten_node lst
        | x -> [ x ] in
      let seqs = List.map aux @@ List.map flatten_node lst in
      Seq (a, List.concat seqs) )
  | x -> x

let descr_to_string descr =
  let node = descr_to_node descr in
  let node = flatten_node node in
  node_to_string node

let n_of_int n =
  match Script_int.is_nat @@ Script_int.of_int n with
  | None -> raise @@ Failure "n_of_int"
  | Some n -> n