open Conv
open Core

exception FromNixError of string

let try_insert_attr x e bs =
  let x = chlist x in
  if Extraction.attr_set_contains x bs then
    raise (FromNixError "Attribute already exists")
  else Extraction.attr_set_insert x e bs

(* Shorthands, minor conversions *)

let mn_singleton_set x e =
  Extraction.(
    EAttr (attr_set_insert (chlist x) (Attr (NONREC, e)) attr_set_empty))

let mn_abs args e =
  List.fold_right args ~init:e ~f:(fun arg e' ->
      Extraction.EAbs (chlist arg, e'))

let mn_lit l = Extraction.ELit l
let mn_int x = mn_lit (Extraction.LitNum (Extraction.NInt x))
let mn_float x = mn_lit (Extraction.LitNum (Extraction.NFloat x))
let mn_bool b = mn_lit (Extraction.LitBool b)
let mn_true = mn_bool true
let mn_false = mn_bool false
let mn_str s = mn_lit (Extraction.LitString (chlist s))
let mn_null = mn_lit Extraction.LitNull
let mn_id x = Extraction.EId (chlist x, None)
let mn_app e1 e2 = Extraction.EApp (e1, e2)
let mn_seq e1 e2 = Extraction.ESeq (Extraction.SHALLOW, e1, e2)
let mn_deep_seq e1 e2 = Extraction.ESeq (Extraction.DEEP, e1, e2)
let mn_list es = Extraction.EList es

let mn_attr (bs : (string * [ `Rec | `Nonrec ] * Extraction.expr) list) =
  Extraction.EAttr
    (List.fold_left bs ~init:Extraction.attr_set_empty ~f:(fun bs' (x, r, e) ->
         let r' =
           match r with `Rec -> Extraction.REC | `Nonrec -> Extraction.NONREC
         in
         Extraction.attr_set_insert (chlist x) (Extraction.Attr (r', e)) bs'))

let mn_with e1 e2 = Extraction.ELetAttr (Extraction.WITH, e1, e2)
let mn_binop op e1 e2 = Extraction.EBinOp (op, e1, e2)
let mn_add e1 e2 = mn_binop Extraction.AddOp e1 e2
let mn_sub e1 e2 = mn_binop Extraction.SubOp e1 e2
let mn_mul e1 e2 = mn_binop Extraction.MulOp e1 e2
let mn_div e1 e2 = mn_binop Extraction.DivOp e1 e2
let mn_bit_and e1 e2 = mn_binop Extraction.AndOp e1 e2
let mn_bit_or e1 e2 = mn_binop Extraction.OrOp e1 e2
let mn_bit_xor e1 e2 = mn_binop Extraction.XOrOp e1 e2
let mn_lt e1 e2 = mn_binop Extraction.LtOp e1 e2
let mn_eq e1 e2 = mn_binop Extraction.EqOp e1 e2
let mn_if e1 e2 e3 = Extraction.EIf (e1, e2, e3)
let mn_delete_attr e1 e2 = mn_binop Extraction.DeleteAttrOp e1 e2
let mn_has_attr e1 e2 = mn_binop Extraction.HasAttrOp e1 e2
let mn_select_attr e1 e2 = mn_binop Extraction.SelectAttrOp e1 e2

let mn_singleton_attr e1 e2 =
  mn_app (mn_binop Extraction.SingletonAttrOp e1 mn_null) e2

let mn_update_attr e1 e2 = mn_binop Extraction.UpdateAttrOp e1 e2
let mn_type_of e = mn_binop Extraction.TypeOfOp e mn_null
let mn_function_args e = mn_binop Extraction.FunctionArgsOp e mn_null
let mn_list_append e1 e2 = mn_binop Extraction.AppendListOp e1 e2
let mn_list_match e = mn_binop Extraction.MatchListOp e mn_null
let mn_string_match e = mn_binop Extraction.MatchStringOp e mn_null
let mn_attr_match e = mn_binop Extraction.MatchAttrOp e mn_null
let mn_ceil e = mn_binop (Extraction.RoundOp Ceil) e mn_null
let mn_nearest_even e = mn_binop (Extraction.RoundOp NearestEven) e mn_null
let mn_floor e = mn_binop (Extraction.RoundOp Floor) e mn_null

(* Macros *)

let mn_cast_bool e = mn_if e mn_true mn_false
let mn_or e1 e2 = mn_if e1 mn_true (mn_cast_bool e2)
let mn_and e1 e2 = mn_if e1 (mn_cast_bool e2) mn_false
let mn_impl e1 e2 = mn_if e1 (mn_cast_bool e2) mn_true
let mn_not e = mn_if e mn_false mn_true
let mn_negate e = mn_sub (mn_int Extraction.Internal.BinNums.Z0) e
let mn_neq e1 e2 = mn_not (mn_eq e2 e1)
let mn_gt e1 e2 = mn_lt e2 e1
let mn_lte e1 e2 = mn_not (mn_gt e1 e2)
let mn_gte e1 e2 = mn_not (mn_lt e1 e2)

(* Macros based on exported functions from the prelude *)

let mnbi_assert e1 e2 = mn_app (mn_app (mn_id "__mn_assert") e1) e2
let mnbi_has_attr e ds = mn_app (mn_app (mn_id "__mn_attr_has") e) (mn_list ds)
let mnbi_select e ds = mn_app (mn_app (mn_id "__mn_attr_select") e) (mn_list ds)

let mnbi_select_or e1 ds e2 =
  mn_app (mn_app (mn_app (mn_id "__mn_attr_selectOr") e1) (mn_list ds)) e2

let mnbi_insert_new e1 e2 e3 =
  mn_app (mn_app (mn_app (mn_id "__mn_attr_insertNew") e1) e2) e3

let is_dynamic_binding (b : Nix.Ast.binding) =
  match b with
  | Nix.Ast.AttrPath ([ Nix.Ast.Val (Nix.Ast.Str (_, [])) ], _)
  | Nix.Ast.Inherit _ ->
      false
  | Nix.Ast.AttrPath ([ _ ], _) -> true
  | _ -> assert false

let has_dynamic_bindings (bs : Nix.Ast.binding list) =
  List.exists bs ~f:is_dynamic_binding

(* Static bindings left, dynamic bindings right *)
let partition_dynamic (bs : Nix.Ast.binding list) :
    Nix.Ast.binding list * Nix.Ast.binding list =
  List.fold_left bs ~init:([], []) ~f:(fun (static, dynamic) b ->
      if is_dynamic_binding b then (static, b :: dynamic)
      else (b :: static, dynamic))

(* Precondition: e must be have been processed by the elaborator. *)
let rec from_nix e =
  match e with
  | Nix.Ast.BinaryOp (op, e1, e2) -> (
      let e1', e2' = (from_nix e1, from_nix e2) in
      match op with
      | Nix.Ast.Plus -> mn_add e1' e2'
      | Nix.Ast.Minus -> mn_sub e1' e2'
      | Nix.Ast.Mult -> mn_mul e1' e2'
      | Nix.Ast.Div -> mn_div e1' e2'
      | Nix.Ast.Gt -> mn_gt e1' e2'
      | Nix.Ast.Lt -> mn_lt e1' e2'
      | Nix.Ast.Lte -> mn_lte e1' e2'
      | Nix.Ast.Gte -> mn_gte e1' e2'
      | Nix.Ast.Eq -> mn_eq e1' e2'
      | Nix.Ast.Neq -> mn_neq e1' e2'
      | Nix.Ast.Or -> mn_or e1' e2'
      | Nix.Ast.And -> mn_and e1' e2'
      | Nix.Ast.Impl -> mn_impl e1' e2'
      | Nix.Ast.Merge -> mn_update_attr e1' e2'
      | Nix.Ast.Concat -> mn_list_append e1' e2')
  | Nix.Ast.UnaryOp (op, e) -> (
      let e = from_nix e in
      match op with Nix.Ast.Negate -> mn_negate e | Nix.Ast.Not -> mn_not e)
  | Nix.Ast.Cond (e1, e2, e3) -> mn_if (from_nix e1) (from_nix e2) (from_nix e3)
  | Nix.Ast.With (e1, e2) -> mn_with (from_nix e1) (from_nix e2)
  | Nix.Ast.Assert (e1, e2) -> mnbi_assert (from_nix e1) (from_nix e2)
  | Nix.Ast.Test (e, ds) -> mnbi_has_attr (from_nix e) List.(ds >>| from_nix)
  | Nix.Ast.SetLet bs ->
      from_nix
        (Nix.Ast.Select
           ( Nix.Ast.Val (Nix.Ast.AttSet (Nix.Ast.Rec, bs)),
             [ Nix.Ast.Val (Nix.Ast.Str ("body", [])) ],
             None ))
  | Nix.Ast.Let (bs, e2) ->
      if has_dynamic_bindings bs then
        raise (FromNixError "Let bindings may not be dynamic");
      let e1 = from_nix (Nix.Ast.Val (Nix.Ast.AttSet (Nix.Ast.Rec, bs))) in
      Extraction.ELetAttr (Extraction.ABS, e1, from_nix e2)
  | Nix.Ast.Val v -> from_nix_val v
  | Nix.Ast.Id x -> mn_id x
  | Nix.Ast.Select (e, parts, md) -> (
      match md with
      | Some d ->
          mnbi_select_or (from_nix e) List.(parts >>| from_nix) (from_nix d)
      | None -> (
          match parts with
          | [ part ] -> mn_select_attr (from_nix e) (from_nix part)
          | _ -> mnbi_select (from_nix e) List.(parts >>| from_nix)))
  | Nix.Ast.Apply (e1, e2) -> mn_app (from_nix e1) (from_nix e2)
  | Nix.Ast.Aquote _ ->
      assert false (* should be gone after processing by elaborator *)

and from_nix_val v =
  match v with
  | Str (s, parts) ->
      let parts = List.(parts >>= fun (e, s) -> [ from_nix e; mn_str s ]) in
      List.fold_left parts ~init:(mn_str s) ~f:mn_add
  | IStr _ -> raise (FromNixError "Indented strings are not supported")
  | Int n -> (
      match Extraction.string_to_Z (chlist n) with
      | Some n -> mn_int n
      | None -> raise (FromNixError "Bad integer literal"))
  | Float n ->
      let n =
        try Float.of_string n
        with Invalid_argument _ -> raise (FromNixError "Bad float literal")
      in
      if Float.(is_nan n || is_inf n) then
        raise (FromNixError "Bad float literal")
      else mn_float (float_to_flocq n)
  | Path _ | SPath _ | HPath _ -> raise (FromNixError "Paths are not supported")
  | Uri s -> mn_str s
  | Lambda (Alias x, e) -> mn_abs [ x ] (from_nix e)
  | Lambda (ParamSet (Some x, fs), e) ->
      from_nix_val
        (Lambda (Alias x, Apply (Val (Lambda (ParamSet (None, fs), e)), Id x)))
  | Lambda (ParamSet (None, (fs, strictness)), e) ->
      let ms =
        List.fold_left fs ~init:Extraction.matcher_empty ~f:(fun ms (x, me) ->
            Extraction.matcher_insert (chlist x) (Option.map ~f:from_nix me) ms)
      in
      Extraction.EAbsMatch
        ( ms,
          (match strictness with Loose -> false | Exact -> true),
          from_nix e )
  | List es -> mn_list (List.map es ~f:from_nix)
  | AttSet (recursivity, bs) ->
      let static, dynamic = partition_dynamic bs
      and recursivity' =
        match recursivity with
        | Nix.Ast.Rec -> Extraction.REC
        | Nix.Ast.Nonrec -> Extraction.NONREC
      in

      let set_no_dyn =
        Extraction.EAttr
          (List.fold_left static ~init:Extraction.attr_set_empty
             ~f:(fun static' bnd ->
               match bnd with
               | Nix.Ast.AttrPath ([ Nix.Ast.Val (Nix.Ast.Str (x, [])) ], e) ->
                   try_insert_attr x
                     (Extraction.Attr (recursivity', from_nix e))
                     static'
               | Nix.Ast.Inherit (None, xs) ->
                   List.fold_left xs ~init:static' ~f:(fun static' x ->
                       match x with
                       | Id x ->
                           try_insert_attr x
                             (Extraction.Attr (Extraction.NONREC, mn_id x))
                             static'
                       | _ -> assert false)
               | Nix.Ast.Inherit (Some e, xs) ->
                   let e = from_nix e in
                   List.fold_left xs ~init:static' ~f:(fun static' x ->
                       match x with
                       | Id x ->
                           try_insert_attr x
                             (Extraction.Attr
                                (recursivity', mn_select_attr e (mn_str x)))
                             static'
                       | _ -> assert false)
               | _ -> assert false))
      in

      List.fold_right dynamic ~init:set_no_dyn ~f:(fun bnd set ->
          match bnd with
          | Nix.Ast.AttrPath ([ d ], e) ->
              mnbi_insert_new set
                (match recursivity with
                | Nix.Ast.Rec ->
                    Extraction.ELetAttr (Extraction.ABS, set_no_dyn, from_nix d)
                | Nix.Ast.Nonrec -> from_nix d)
                (match recursivity with
                | Nix.Ast.Rec ->
                    Extraction.ELetAttr (Extraction.ABS, set_no_dyn, from_nix e)
                | Nix.Ast.Nonrec -> from_nix e)
          | _ -> assert false)