open Core
open Types

(* The Nix elaborator does a few things:
   - Attribute paths are transformed into a simple list of expressions:
     + Simple identifiers are rewritten to string values
     + Antiquotations are rewritten to their component expressions
     + Anything else, that is not a string value, is rejected
       and raises an exception
   - In 'inherit (...) x1 ... xn', x1 ... xn are checked for 'reasonably' being
     identifiers, i.e., being one of x, "x" and ${"x"}.
   - Nested attribute paths are unfolded and attribute sets are merged where
     possible. (Where we mean 'what Nix does' with 'where possible'; see the
     comment at the respective function.)
   - Paths are turned into strings and made absolute w.r.t. the current
     working directory.
   - Indented strings are converted to their 'normal' counterpart. *)

exception ElaborateError of string

type attr_set = recursivity * binding list

let set_expr (r, bs) = Val (AttSet (r, bs))
let get_id = function Id x -> x | _ -> assert false

let rec update_bnd (bs : binding list) (x : string) ~(f : expr option -> expr) =
  match bs with
  | [] -> [ AttrPath ([ Val (Str (x, [])) ], f None) ]
  | AttrPath ([ Val (Str (y, [])) ], e) :: s' when String.(x = y) ->
      AttrPath ([ Val (Str (y, [])) ], f (Some e)) :: s'
  | Inherit (_, ids) :: _
    when List.exists ids ~f:(fun e -> String.(get_id e = x)) ->
      raise (ElaborateError "Cannot update inherit")
  | bnd :: s' -> bnd :: update_bnd s' x ~f

let set_update_bnd (r, bs) x ~f = (r, update_bnd bs x ~f)

let rec has_bnd (bs : binding list) (x : string) : bool =
  match bs with
  | [] -> false
  | AttrPath ([ Val (Str (y, [])) ], _) :: _ when String.(x = y) -> true
  | Inherit (_, ids) :: _
    when List.exists ids ~f:(fun e -> String.(get_id e = x)) ->
      true
  | _ :: bs' -> has_bnd bs' x

let merge_bnds bs1 bs2 : binding list =
  List.fold_left bs2 ~init:bs1 ~f:(fun bs1' b2 ->
      match b2 with
      | AttrPath ([ Val (Str (x, [])) ], e) ->
          update_bnd bs1' x ~f:(function
            | Some _ -> raise (ElaborateError "Duplicated attribute")
            | None -> e)
      | AttrPath ([ d ], e) -> AttrPath ([ d ], e) :: bs1'
      | Inherit (md, xs) ->
          if List.for_all xs ~f:(fun e -> not (has_bnd bs1' (get_id e))) then
            Inherit (md, xs) :: bs1'
          else raise (ElaborateError "Duplicated attribute")
      | _ -> assert false)

(* This function intentionally clobbers recursivity, because that is the way
   that Nix likes to handle attribute insertion. See
     (1) https://github.com/NixOS/nix/issues/9020
     (2) https://github.com/NixOS/nix/issues/11268
     (3) https://github.com/NixOS/nix/pull/11294 *)
let rec insert (bs : binding list) (path : expr list) (e : expr) =
  match path with
  | [] -> raise (ElaborateError "Cannot insert attribute with empty path")
  | [ Val (Str (x, [])) ] ->
      update_bnd bs x ~f:(function
        | None -> e
        | Some (Val (AttSet (r1, bs1))) -> (
            match e with
            | Val (AttSet (_, bs2)) -> set_expr (r1, merge_bnds bs1 bs2)
            | _ -> raise (ElaborateError "Duplicated attribute"))
        | _ -> raise (ElaborateError "Duplicated attribute"))
  | Val (Str (x, [])) :: rest ->
      update_bnd bs x ~f:(function
        | Some (Val (AttSet (r, bs))) -> Val (AttSet (r, insert bs rest e))
        | Some _ -> raise (ElaborateError "Duplicated attribute")
        | None -> Val (AttSet (Nonrec, insert [] rest e)))
  | [ part ] -> AttrPath ([ part ], e) :: bs
  | part :: rest ->
      AttrPath ([ part ], Val (AttSet (Nonrec, insert [] rest e))) :: bs

let insert_inherit (bs : binding list) (from : expr option) (es : expr list) =
  if List.for_all es ~f:(fun e -> not (has_bnd bs (get_id e))) then
    Inherit (from, es) :: bs
  else raise (ElaborateError "Duplicated attribute")

let simplify_path_component = function
  | Id x -> Val (Str (x, []))
  | Val (Str (s, ess)) -> Val (Str (s, ess))
  | Aquote e -> e
  | _ -> raise (ElaborateError "Unexpected path component")

let simplify_path = List.map ~f:simplify_path_component

let simplify_bnd_paths =
  List.map ~f:(fun bnd ->
      match bnd with
      | AttrPath (path, e) -> AttrPath (simplify_path path, e)
      | Inherit (me, xs) -> Inherit (me, xs))

(* Law: concat_lines ∘ split_lines = id *)

let rec split_lines s =
  match String.lsplit2 s ~on:'\n' with
  | Some (s1, s2) -> s1 :: split_lines s2
  | None -> [ s ]

let rec concat_lines = function
  | [] -> ""
  | [ x ] -> x
  | x :: xs -> x ^ "\n" ^ concat_lines xs

let map_tail ~f = function [] -> [] | x :: xs -> x :: List.map ~f xs

let unindent n s ~skip_first_line =
  let map_op ~f = if skip_first_line then map_tail ~f else List.map ~f in
  split_lines s
  |> map_op ~f:(fun line ->
         let expected_prefix = String.make n ' ' in
         String.chop_prefix_if_exists ~prefix:expected_prefix line)
  |> concat_lines

let is_spaces l = String.(strip l ~drop:(Char.( = ) ' ') |> is_empty)

let drop_first_empty_line s =
  match String.lsplit2 s ~on:'\n' with
  | Some (l, s') when is_spaces l -> s'
  | _ -> s

let rec process ?(dir = None) = function
  | BinaryOp (op, e1, e2) -> BinaryOp (op, process ~dir e1, process ~dir e2)
  | UnaryOp (op, e) -> UnaryOp (op, process ~dir e)
  | Cond (e1, e2, e3) -> Cond (process ~dir e1, process ~dir e2, process ~dir e3)
  | With (e1, e2) -> With (process ~dir e1, process ~dir e2)
  | Assert (e1, e2) -> Assert (process ~dir e1, process ~dir e2)
  | Test (e1, es) ->
      Test (process ~dir e1, List.(simplify_path es >>| process ~dir))
  | SetLet bs -> SetLet (process_bnds ~dir bs)
  | Let (bs, e) -> Let (process_bnds ~dir bs, process ~dir e)
  | Val v -> Val (process_val ~dir v)
  | Id x -> Id x
  | Select (e, es, me) ->
      Select
        ( process ~dir e,
          List.(simplify_path es >>| process ~dir),
          Option.(me >>| process ~dir) )
  | Apply (e1, e2) -> Apply (process ~dir e1, process ~dir e2)
  | Aquote e -> Aquote (process ~dir e)

and process_val ~dir = function
  | Str (s, ess) -> Str (s, List.(ess >>| fun (e, s) -> (process ~dir e, s)))
  | IStr (n, s, ess) ->
      let s' = drop_first_empty_line (unindent n s ~skip_first_line:false)
      and ess' =
        List.map ess ~f:(fun (e, s) ->
            (process ~dir e, unindent n s ~skip_first_line:true))
      in
      Str (s', ess')
  | Lambda (p, e) -> Lambda (process_pattern ~dir p, process ~dir e)
  | List es -> List List.(es >>| process ~dir)
  | AttSet (r, bs) -> AttSet (r, process_bnds ~dir bs)
  | Path p -> (
      if Filename.is_absolute p then Str (p, [])
      else
        match dir with
        | Some dir when Filename.is_absolute dir ->
            Str (Filename.concat dir p, [])
        | Some _ ->
            raise
              (ElaborateError "Provided directory should be an absolute path")
        | None -> raise (ElaborateError "Do not know how to resolve path"))
  | v -> v

and process_bnds ~dir bs =
  bs
  |> List.map ~f:(function
       | AttrPath (es, e) ->
           AttrPath (List.(es >>| process ~dir), process ~dir e)
       | Inherit (me, xs) ->
           Inherit (Option.(me >>| process ~dir), process_inherit_ids xs))
  |> simplify_bnd_paths
  |> List.fold ~init:[] ~f:(fun bs' bnd ->
         match bnd with
         | AttrPath (path, e) -> insert bs' path e
         | Inherit (from, es) -> insert_inherit bs' from es)

and process_inherit_ids =
  List.map ~f:(function
    | Id x | Val (Str (x, [])) | Aquote (Val (Str (x, []))) -> Id x
    | _ -> raise (ElaborateError "Unexpected expression in inherit"))

and process_pattern ~dir = function
  | Alias x -> Alias x
  | ParamSet (mx, (ps, k)) -> ParamSet (mx, (process_param_set ~dir mx ps, k))

and process_param_set ~dir ?(seen = String.Set.empty) mx ps =
  match ps with
  | [] -> []
  | (y, me) :: ps' ->
      if Set.mem seen y || Option.mem mx y ~equal:String.( = ) then
        raise (ElaborateError "Duplicated function argument")
      else
        (y, Option.(me >>| process ~dir))
        :: process_param_set ~dir mx ps' ~seen:(Set.add seen y)