open Trace
open Mini_c
open Michelson
open Memory_proto_alpha.Protocol.Script_ir_translator
open Operators.Compiler

let get_operator : string -> type_value -> expression list -> predicate result = fun s ty lst ->
  match Map.String.find_opt s Operators.Compiler.operators with
  | Some x -> ok x
  | None -> (
      match s with
      | "NONE" -> (
          let%bind ty' = Mini_c.get_t_option ty in
          let%bind m_ty = Compiler_type.type_ ty' in
          ok @@ simple_constant @@ prim ~children:[m_ty] I_NONE
        )
      | "NIL" -> (
          let%bind ty' = Mini_c.get_t_list ty in
          let%bind m_ty = Compiler_type.type_ ty' in
          ok @@ simple_unary @@ prim ~children:[m_ty] I_NIL
        )
      | "SET_EMPTY" -> (
          let%bind ty' = Mini_c.get_t_set ty in
          let%bind m_ty = Compiler_type.type_ ty' in
          ok @@ simple_constant @@ prim ~children:[m_ty] I_EMPTY_SET
        )
      | "UNPACK" -> (
          let%bind ty' = Mini_c.get_t_option ty in
          let%bind m_ty = Compiler_type.type_ ty' in
          ok @@ simple_unary @@ prim ~children:[m_ty] I_UNPACK
        )
      | "MAP_REMOVE" ->
          let%bind v = match lst with
            | [ _ ; expr ] ->
                let%bind (_, v) = Mini_c.Combinators.(get_t_map (Expression.get_type expr)) in
                ok v
            | _ -> simple_fail "mini_c . MAP_REMOVE" in
          let%bind v_ty = Compiler_type.type_ v in
          ok @@ simple_binary @@ seq [dip (i_none v_ty) ; prim I_UPDATE ]
      | "LEFT" ->
          let%bind r = match lst with
            | [ _ ] -> get_t_right ty
            | _ -> simple_fail "mini_c . LEFT" in
          let%bind r_ty = Compiler_type.type_ r in
          ok @@ simple_unary @@ prim ~children:[r_ty] I_LEFT
      | "RIGHT" ->
          let%bind l = match lst with
            | [ _ ] -> get_t_left ty
            | _ -> simple_fail "mini_c . RIGHT" in
          let%bind l_ty = Compiler_type.type_ l in
          ok @@ simple_unary @@ prim ~children:[l_ty] I_RIGHT
      | "CONTRACT" ->
          let%bind r = match lst with
            | [ _ ] -> get_t_contract ty
            | _ -> simple_fail "mini_c . CONTRACT" in
          let%bind r_ty = Compiler_type.type_ r in
          ok @@ simple_unary @@ seq [
            prim ~children:[r_ty] I_CONTRACT ;
            i_assert_some_msg (i_push_string "bad address for get_contract") ;
          ]
      | x -> simple_fail ("predicate \"" ^ x ^ "\" doesn't exist")
    )

let rec translate_value (v:value) ty : michelson result = match v with
  | D_bool b -> ok @@ prim (if b then D_True else D_False)
  | D_int n -> ok @@ int (Z.of_int n)
  | D_nat n -> ok @@ int (Z.of_int n)
  | D_timestamp n -> ok @@ int (Z.of_int n)
  | D_tez n -> ok @@ int (Z.of_int n)
  | D_string s -> ok @@ string s
  | D_bytes s -> ok @@ bytes (Tezos_stdlib.MBytes.of_bytes s)
  | D_unit -> ok @@ prim D_Unit
  | D_pair (a, b) -> (
      let%bind (a_ty , b_ty) = get_t_pair ty in
      let%bind a = translate_value a a_ty in
      let%bind b = translate_value b b_ty in
      ok @@ prim ~children:[a;b] D_Pair
    )
  | D_left a -> (
      let%bind (a_ty , _) = get_t_or ty in
      let%bind a' = translate_value a a_ty in
      ok @@ prim ~children:[a'] D_Left
    )
  | D_right b -> (
      let%bind (_ , b_ty) = get_t_or ty in
      let%bind b' = translate_value b b_ty in
      ok @@ prim ~children:[b'] D_Right
    )
  | D_function func -> (
      match ty with
      | T_function (in_ty , _) -> translate_function_body func [] in_ty
      | _ -> simple_fail "expected function type"
    )
  | D_none -> ok @@ prim D_None
  | D_some s ->
      let%bind s' = translate_value s ty in
      ok @@ prim ~children:[s'] D_Some
  | D_map lst -> (
      let%bind (k_ty , v_ty) = get_t_map ty in
      let%bind lst' =
        let aux (k , v) = bind_pair (translate_value k k_ty , translate_value v v_ty) in
        bind_map_list aux lst in
      let sorted = List.sort (fun (x , _) (y , _) -> compare x y) lst' in
      let aux (a, b) = prim ~children:[a;b] D_Elt in
      ok @@ seq @@ List.map aux sorted
    )
  | D_list lst -> (
      let%bind e_ty = get_t_list ty in
      let%bind lst' = bind_map_list (fun x -> translate_value x e_ty) lst in
      ok @@ seq lst'
    )
  | D_set lst -> (
      let%bind e_ty = get_t_set ty in
      let%bind lst' = bind_map_list (fun x -> translate_value x e_ty) lst in
      let sorted = List.sort compare lst' in
      ok @@ seq sorted
    )
  | D_operation _ ->
      simple_fail "can't compile an operation"

and translate_expression (expr:expression) (env:environment) : michelson result =
  let (expr' , ty) = Combinators.Expression.(get_content expr , get_type expr) in
  let error_message () =
    Format.asprintf  "\n- expr: %a\n- type: %a\n" PP.expression expr PP.type_ ty
  in
  let return code = ok code in

  trace (error (thunk "compiling expression") error_message) @@
  match expr' with
  | E_skip -> return @@ i_push_unit
  | E_literal v ->
      let%bind v = translate_value v ty in
      let%bind t = Compiler_type.type_ ty in
      return @@ i_push t v
  | E_closure anon -> (
      match ty with
      | T_deep_closure (small_env , input_ty , output_ty) -> (
          let selector = List.map fst small_env in
          let%bind closure_pack_code = Compiler_environment.pack_closure env selector in
          let%bind lambda_ty = Compiler_type.lambda_closure (small_env , input_ty , output_ty) in
          let%bind lambda_body_code = translate_function_body anon small_env input_ty in
          return @@ seq [
            closure_pack_code ;
            i_push lambda_ty lambda_body_code ;
            i_pair ;
          ]
        )
      | _ -> simple_fail "expected closure type"
    )
  | E_application (f , arg) -> (
      match Combinators.Expression.get_type f with
      | T_function _ -> (
          trace (simple_error "Compiling quote application") @@
          let%bind f = translate_expression f env in
          let%bind arg = translate_expression arg env in
          return @@ seq [
            arg ;
            dip f ;
            prim I_EXEC ;
          ]
        )
      | T_deep_closure (_ , _ , _) -> (
          let%bind f_code = translate_expression f env in
          let%bind arg_code = translate_expression arg env in
          return @@ seq [
            arg_code ;
            dip (seq [ f_code ; i_unpair ; i_swap ]) ; i_pair ;
            prim I_EXEC ;
          ]
        )
      | _ -> simple_fail "E_applicationing something not appliable"
    )
  | E_variable x ->
    let%bind code = Compiler_environment.get env x in
    return code
  | E_sequence (a , b) -> (
    let%bind a' = translate_expression a env in
    let%bind b' = translate_expression b env in
    return @@ seq [
      a' ;
      i_drop ;
      b' ;
    ]
  )
  | E_constant(str, lst) ->
      let module L = Logger.Stateful() in
      let%bind pre_code =
        let aux code expr =
          let%bind expr_code = translate_expression expr env in
          L.log @@ Format.asprintf "\n%a -> %a in %a\n"
            PP.expression expr
            Michelson.pp expr_code
            PP.environment env ;
          ok (seq [ expr_code ; dip code ]) in
        bind_fold_right_list aux (seq []) lst in
      let%bind predicate = get_operator str ty lst in
      let%bind code = match (predicate, List.length lst) with
        | Constant c, 0 -> ok @@ seq [
            pre_code ;
            c ;
          ]
        | Unary f, 1 -> ok @@ seq [
            pre_code ;
            f ;
          ]
        | Binary f, 2 -> ok @@ seq [
            pre_code ;
            f ;
          ]
        | Ternary f, 3 -> ok @@ seq [
            pre_code ;
            f ;
          ]
        | _ -> simple_fail "bad arity"
      in
      let error =
        let title () = "error compiling constant" in
        let content () = L.get () in
        error title content in
      trace error @@
      return code
  | E_make_empty_map sd ->
      let%bind (src, dst) = bind_map_pair Compiler_type.type_ sd in
      return @@ i_empty_map src dst
  | E_make_empty_list t ->
      let%bind t' = Compiler_type.type_ t in
      return @@ i_nil t'
  | E_make_empty_set t ->
      let%bind t' = Compiler_type.type_ t in
      return @@ i_empty_set t'
  | E_make_none o ->
      let%bind o' = Compiler_type.type_ o in
      return @@ i_none o'
  | E_if_bool (c, a, b) -> (
      let%bind c' = translate_expression c env in
      let%bind a' = translate_expression a env in
      let%bind b' = translate_expression b env in
      let%bind code = ok (seq [
          c' ;
          i_if a' b' ;
        ]) in
      return code
    )
  | E_if_none (c, n, (ntv , s)) -> (
      let%bind c' = translate_expression c env in
      let%bind n' = translate_expression n env in
      let s_env = Environment.add ntv env in
      let%bind s' = translate_expression s s_env in
      let%bind code = ok (seq [
          c' ;
          i_if_none n' (seq [
              s' ;
              dip i_drop ;
            ])
          ;
        ]) in
      return code
    )
  | E_if_left (c, (l_ntv , l), (r_ntv , r)) -> (
      let%bind c' = translate_expression c env in
      let l_env = Environment.add l_ntv env in
      let%bind l' = translate_expression l l_env in
      let r_env = Environment.add r_ntv env in
      let%bind r' = translate_expression r r_env in
      let%bind code = ok (seq [
          c' ;
          i_if_left (seq [
              l' ;
              i_comment "restrict left" ;
              dip i_drop ;
            ]) (seq [
              r' ;
              i_comment "restrict right" ;
              dip i_drop ;
            ])
          ;
        ]) in
      return code
    )
  | E_let_in (v , expr , body) -> (
      let%bind expr' = translate_expression expr env in
      let%bind body' = translate_expression body (Environment.add v env) in
      let%bind code = ok (seq [
          expr' ;
          body' ;
          i_comment "restrict let" ;
          dip i_drop ;
        ]) in
      return code
    )
  | E_iterator (name , (v , body) , expr) -> (
      let%bind expr' = translate_expression expr env in
      let%bind body' = translate_expression body (Environment.add v env) in
      match name with
      | "ITER" -> (
          let%bind code = ok (seq [
              expr' ;
              i_iter (seq [body' ; i_drop ; i_drop]) ;
              i_push_unit ;
            ]) in
          return code
        )
      | "MAP" -> (
          let%bind code = ok (seq [
              expr' ;
              i_map (seq [body' ; dip i_drop]) ;
            ]) in
          return code
        )
      | s -> (
          let error = error (thunk "bad iterator") (thunk s) in
          fail error
        )
    )
  | E_assignment (name , lrs , expr) -> (
      let%bind expr' = translate_expression expr env in
      let%bind get_code = Compiler_environment.get env name in
      let modify_code =
        let aux acc step = match step with
          | `Left -> seq [dip i_unpair ; acc ; i_pair]
          | `Right -> seq [dip i_unpiar ; acc ; i_piar]
        in
        let init = dip i_drop in
        List.fold_right' aux init lrs
      in
      let%bind set_code = Compiler_environment.set env name in
      let error =
        let title () = "michelson type-checking patch" in
        let content () =
          let aux ppf = function
            | `Left -> Format.fprintf ppf "left"
            | `Right -> Format.fprintf ppf "right" in
          Format.asprintf "Sub path: %a\n"
            PP_helpers.(list_sep aux (const " , ")) lrs
        in
        error title content in
      trace error @@
      return @@ seq [
        i_comment "assign: start # env" ;
        expr' ;
        i_comment "assign: compute rhs # rhs : env" ;
        dip get_code ;
        i_comment "assign: get name # rhs : name : env" ;
        modify_code ;
        i_comment "assign: modify code # name+rhs : env" ;
        set_code ;
        i_comment "assign: set new # new_env" ;
        i_push_unit ;
      ]
    )
  | E_while (expr , block) -> (
      let%bind expr' = translate_expression expr env in
      let%bind block' = translate_expression block env in
      return @@ seq [
        expr' ;
        prim ~children:[seq [
            block' ;
            i_drop ;
            expr']] I_LOOP ;
        i_push_unit ;
      ]
    )

and translate_function_body ({result ; binder} : anon_function) lst input : michelson result =
  let pre_env = Environment.of_list lst in
  let env = Environment.(add (binder , input) pre_env) in
  let%bind expr_code = translate_expression result env in
  let%bind unpack_closure_code = Compiler_environment.unpack_closure pre_env in
  let code = seq [
      i_comment "unpack closure env" ;
      unpack_closure_code ;
      i_comment "function result" ;
      expr_code ;
      i_comment "remove env" ;
      dip i_drop ;
      seq (List.map (Function.constant (dip i_drop)) lst) ;
    ] in

  ok code

type compiled_program = {
  input : ex_ty ;
  output : ex_ty ;
  body : michelson ;
}

let get_main : program -> string -> (anon_function * _) result = fun p entry ->
  let is_main (((name , expr), _):toplevel_statement) =
    match Combinators.Expression.(get_content expr , get_type expr)with
    | (E_literal (D_function content) , T_function ty)
      when name = entry ->
        Some (content , ty)
    | _ ->  None
  in
  let%bind main =
    trace_option (simple_error "no functional entry") @@
    List.find_map is_main p
  in
  ok main

let translate_program (p:program) (entry:string) : compiled_program result =
  let%bind (main , (input , output)) = get_main p entry in
  let%bind body = translate_function_body main [] input in
  let%bind input = Compiler_type.Ty.type_ input in
  let%bind output = Compiler_type.Ty.type_ output in
  ok ({input;output;body}:compiled_program)

let translate_entry (p:anon_function) ty : compiled_program result =
  let (input , output) = ty in
  let%bind body =
    trace (simple_error "compile entry body") @@
    translate_function_body p [] input in
  let%bind input = Compiler_type.Ty.type_ input in
  let%bind output = Compiler_type.Ty.type_ output in
  ok ({input;output;body}:compiled_program)

module Errors = struct
  let corner_case ~loc message =
    let title () = "corner case" in
    let content () = "we don't have a good error message for this case. we are
striving find ways to better report them and find the use-cases that generate
them. please report this to the developers." in
    let data = [
      ("location" , fun () -> loc) ;
      ("message" , fun () -> message) ;
    ] in
    error ~data title content
end
open Errors

let translate_contract : anon_function -> _ -> michelson result = fun f ty ->
  let%bind compiled_program =
    trace_strong (corner_case ~loc:__LOC__ "compiling") @@
    translate_entry f ty in
  let%bind (param_ty , storage_ty) = Combinators.get_t_pair (fst ty) in
  let%bind param_michelson = Compiler_type.type_ param_ty in
  let%bind storage_michelson = Compiler_type.type_ storage_ty in
  let contract = Michelson.contract param_michelson storage_michelson compiled_program.body in
  ok contract