(** Forms → AST. Recognises special forms, desugars sugar, reports malformed
    syntax with the location of the offending form.

    Everything not recognised here is a call, which is how a Lisp should work:
    [at], [len], [push], [abort] and the rest are ordinary functions resolved
    by the checker, not syntax. *)

open Form

let fail (f : Form.t) fmt = Loc.fail f.loc fmt

let sym (f : Form.t) =
  match f.v with
  | Sym s -> s
  | _ -> fail f "expected a name, found %s" (Form.to_string f)

(* Primitive type names are lowercase but concrete; every other lowercase name
   in type position is a type variable (plan.org, Types). *)
let primitives =
  [ "i8"; "i16"; "i32"; "i64"; "u8"; "u16"; "u32"; "u64";
    "f32"; "f64"; "bool"; "string"; "Unit"; "Never" ]

let is_primitive s = List.mem s primitives

module Names = Set.Make (String)

(* Type constructors the language provides. User types are collected by a
   pre-pass over the file's declarations — see [program]. *)
let builtin_types =
  Names.of_list (primitives @ [ "Ptr"; "Option"; "Result"; "Vec"; "Map";
                                "Handle"; "Fn" ])

(* ── Type expressions ──────────────────────────────────────────────── *)

let rec texpr (f : Form.t) : Ast.texpr =
  let mk t = { Ast.t; tloc = f.loc } in
  match f.v with
  | Sym s -> mk (Ast.Tname s)
  | Vec [ elem ] -> mk (Ast.Tslice (texpr elem))
  | Vec [ n; elem ] -> mk (Ast.Tarray (len n, texpr elem))
  | Vec _ ->
    fail f "a type in brackets is [T] for a slice or [n T] for a fixed array"
  | Map [ k; v ] -> mk (Ast.Tmap (texpr k, texpr v))
  | Map _ -> fail f "a map type is {K V}"
  | List ({ v = Sym "Fn"; _ } :: rest) ->
    (match rest with
     | [ { v = Vec params; _ }; ret ] ->
       mk (Ast.Tfn (List.map texpr params, texpr ret))
     | _ -> fail f "a function type is (Fn [T ...] R)")
  | List ({ v = Sym name; _ } :: args) when args <> [] ->
    mk (Ast.Tapp (name, List.map texpr args))
  | _ -> fail f "expected a type, found %s" (Form.to_string f)

and len (f : Form.t) : Ast.len =
  match f.v with
  | Int n -> Ast.Lint n
  | Sym s -> Ast.Lname s
  | _ -> fail f "an array length is an integer or a constant's name"

(* Inline name/type pairs: [x i32  y f32] — as in let, defstruct and defn. *)
let rec fields (f : Form.t) (items : Form.t list) : Ast.field list =
  match items with
  | [] -> []
  | name :: ty :: rest ->
    { Ast.fname = sym name; fty = texpr ty; floc = name.loc } :: fields f rest
  | [ odd ] ->
    Loc.fail odd.loc "field %s has no type — these come in name/type pairs"
      (Form.to_string odd)

(* ── Expressions ───────────────────────────────────────────────────── *)

let rec expr (f : Form.t) : Ast.expr =
  let mk e = { Ast.e; loc = f.loc } in
  match f.v with
  | Int i -> mk (Ast.Int i)
  | Float x -> mk (Ast.Float x)
  | Byte b -> mk (Ast.Byte b)
  | Str s -> mk (Ast.Str s)
  | Kw k -> mk (Ast.Kw k)
  | Sym s -> mk (Ast.Var s)
  (* In value position brackets are a fixed-array literal; in type position
     they are a slice or array type. Position disambiguates, as with {}. *)
  | Vec items -> mk (Ast.Arr (List.map expr items))
  | Map _ -> fail f "a bare map is not an expression; write (Type {:field v})"
  | List [] -> fail f "() is not an expression"
  | List (head :: args) -> form f mk head args

and form f mk (head : Form.t) (args : Form.t list) : Ast.expr =
  match head.v with
  (* ── quote ─────────────────────────────────────────────────────── *)
  | Sym "quote" ->
    (match args with
     | [ { v = Sym s; _ } ] -> mk (Ast.Quote s)
     | _ -> fail f "quote takes one symbol")

  (* ── sequencing and binding ────────────────────────────────────── *)
  | Sym "do" -> mk (Ast.Do (List.map expr args))

  | Sym "let" ->
    (match args with
     | { v = Vec bs; _ } :: body -> mk (Ast.Let (bindings f bs, body_of body))
     | _ -> fail f "let is (let [name value ...] body ...)")

  (* ── conditionals ──────────────────────────────────────────────── *)
  | Sym "if" ->
    (match args with
     | [ c; t ] -> mk (Ast.If (expr c, expr t, None))
     | [ c; t; e ] -> mk (Ast.If (expr c, expr t, Some (expr e)))
     | _ -> fail f "if is (if test then) or (if test then else)")

  (* Sugar, desugared here: special forms until macros land at milestone 5. *)
  | Sym "when" ->
    (match args with
     | c :: body when body <> [] ->
       mk (Ast.If (expr c, { Ast.e = Ast.Do (body_of body); loc = f.loc }, None))
     | _ -> fail f "when is (when test body ...)")

  | Sym "unless" ->
    (match args with
     | c :: body when body <> [] ->
       let neg = { Ast.e = Ast.Call ({ Ast.e = Ast.Var "not"; loc = head.loc },
                                     [ expr c ]); loc = f.loc } in
       mk (Ast.If (neg, { Ast.e = Ast.Do (body_of body); loc = f.loc }, None))
     | _ -> fail f "unless is (unless test body ...)")

  | Sym "cond" -> cond f args

  (* Short-circuiting, so they cannot be ordinary calls. *)
  | Sym "and" -> shortcircuit f args ~is_and:true
  | Sym "or" -> shortcircuit f args ~is_and:false

  (* ── loops ─────────────────────────────────────────────────────── *)
  | Sym "while" ->
    (match args with
     | c :: body -> mk (Ast.While (expr c, body_of body))
     | [] -> fail f "while is (while test body ...)")

  | Sym "until" ->
    (match args with
     | c :: body ->
       let neg = { Ast.e = Ast.Call ({ Ast.e = Ast.Var "not"; loc = head.loc },
                                     [ expr c ]); loc = f.loc } in
       mk (Ast.While (neg, body_of body))
     | [] -> fail f "until is (until test body ...)")

  (* ── control ───────────────────────────────────────────────────── *)
  | Sym "return" ->
    (match args with
     | [] -> mk (Ast.Return None)
     | [ v ] -> mk (Ast.Return (Some (expr v)))
     | _ -> fail f "return takes at most one value")

  | Sym "set" ->
    (match args with
     | [ target; value ] -> mk (Ast.Set (place target, expr value))
     | _ -> fail f "set is (set place value)")

  | Sym "match" ->
    (match args with
     | scrutinee :: rest -> mk (Ast.Match (expr scrutinee, arms f rest))
     | [] -> fail f "match is (match value pattern body ...)")

  (* ── binding and control: never a call ─────────────────────────── *)
  (* A form that binds a name or alters control flow cannot fall through to
     Call — it would parse cleanly and mean the wrong thing, silently. *)
  | Sym "fn" ->
    (match args with
     | { v = Vec ps; _ } :: body when body <> [] ->
       mk (Ast.Fn (List.map sym ps, body_of body))
     | _ -> fail f "fn is (fn [param ...] body ...)")

  | Sym "dotimes" ->
    (match args with
     | { v = Vec [ n; count ]; _ } :: body ->
       mk (Ast.Dotimes (sym n, expr count, body_of body))
     | _ -> fail f "dotimes is (dotimes [name count] body ...)")

  | Sym "defer" ->
    (match args with
     | [] -> fail f "defer is (defer body ...)"
     | body -> mk (Ast.Defer (body_of body)))

  | Sym "some" ->
    (match args with
     | [ v ] -> mk (Ast.Unwrap (Ast.Usome, expr v))
     | _ -> fail f "some is (some option-value)")

  | Sym "try" ->
    (match args with
     | [ v ] -> mk (Ast.Unwrap (Ast.Utry, expr v))
     | _ -> fail f "try is (try result-value)")

  (* (signal c) : Unit, always. When every applicable handler returns normally
     the signalling function simply carries on, and with no handler at all it is
     a no-op — spec-conditions.md §1 and §2. *)
  | Sym "signal" ->
    (match args with
     | [ c ] -> mk (Ast.Signal (expr c))
     | _ -> fail f "signal is (signal condition)")

  (* (handler-bind [(Type [c] body ...) ...] body ...)

     A clause names a condition type, binds the condition, and runs for effect;
     matching is by type, since there is no condition hierarchy. *)
  | Sym "handler-bind" ->
    let clauses, body =
      match args with
      | { v = Vec clauses; _ } :: body when body <> [] -> (clauses, body)
      | _ ->
        fail f "handler-bind is (handler-bind [(Type [name] body ...) ...] body ...)"
    in
    let clause (c : Form.t) =
      match c.Form.v with
      | Form.List (ty :: { v = Form.Vec [ { v = Form.Sym n; _ } ]; _ } :: cbody)
        when cbody <> [] ->
        { Ast.hty = texpr ty; hname = n; hbody = List.map expr cbody;
          hloc = c.Form.loc }
      | _ ->
        fail c "a handler-bind clause is (Type [name] body ...)"
    in
    mk (Ast.HandlerBind (List.map clause clauses, body_of body))

  (* Recognised, deliberately unimplemented. Rejected rather than left to fall
     through to Call, where they would parse and mean nothing. *)
  | Sym ("handler-case" | "restart-case" | "invoke-restart"
        | "errdefer" | "with-allocator" | "loop" | "recur"
        | "defmacro" | "await" as name) ->
    fail f "%s is not implemented yet (see the build sequence in plan.org)" name

  (* ── field access: (.pos c) ────────────────────────────────────── *)
  | Sym s when String.length s > 1 && s.[0] = '.' ->
    let field = String.sub s 1 (String.length s - 1) in
    (match args with
     | [ target ] -> mk (Ast.Field (expr target, field))
     | _ -> fail f "field access is (.%s value)" field)

  (* ── struct literal: (Cursor {:src s :pos 0}) ───────────────────── *)
  | Sym name when args <> [] && is_map (List.hd args) ->
    (match args with
     | [ { v = Map kvs; _ } ] -> mk (Ast.Struct (name, struct_fields f kvs))
     | _ -> fail f "a struct literal is (%s {:field value ...})" name)

  (* ── anything else is a call ────────────────────────────────────── *)
  | _ -> mk (Ast.Call (expr head, List.map expr args))

and is_map (f : Form.t) = match f.v with Map _ -> true | _ -> false

and body_of (items : Form.t list) : Ast.expr list = List.map expr items

and bindings f (items : Form.t list) : Ast.binding list =
  (* [name value ...] and [name Type value ...] both read; a type is a form
     that is not a value position — disambiguated by pair vs triple is
     ambiguous, so let requires (let [name value]) and types are inferred.
     Annotated locals are not needed by any acceptance program. *)
  let rec go = function
    | [] -> []
    | name :: value :: rest ->
      { Ast.bname = sym name; bty = None; bval = expr value; bloc = name.loc }
      :: go rest
    | [ odd ] ->
      Loc.fail odd.loc "binding %s has no value — let takes name/value pairs"
        (Form.to_string odd)
  in
  if items = [] then Loc.fail f.loc "let needs at least one binding" else go items

and struct_fields f (items : Form.t list) : (string * Ast.expr) list =
  let rec go = function
    | [] -> []
    | { v = Kw k; _ } :: value :: rest -> (k, expr value) :: go rest
    | other :: _ :: _ ->
      Loc.fail other.loc "expected :field, found %s" (Form.to_string other)
    | [ odd ] -> Loc.fail odd.loc "field %s has no value" (Form.to_string odd)
  in
  ignore f; go items

and cond f (args : Form.t list) : Ast.expr =
  let rec go = function
    | [] -> { Ast.e = Ast.Do []; loc = f.loc }   (* no clause matched: Unit *)
    | { v = Kw "else"; _ } :: body :: _ -> expr body
    | test :: body :: rest ->
      { Ast.e = Ast.If (expr test, expr body, Some (go rest)); loc = f.loc }
    | [ odd ] -> Loc.fail odd.loc "cond clause %s has no body"
                   (Form.to_string odd)
  in
  if args = [] then Loc.fail f.loc "cond needs at least one clause" else go args

and shortcircuit f (args : Form.t list) ~is_and : Ast.expr =
  let mk e = { Ast.e; loc = f.loc } in
  let rec go = function
    | [] -> mk (Ast.Var (if is_and then "true" else "false"))
    | [ last ] -> expr last
    | x :: rest ->
      if is_and then mk (Ast.If (expr x, go rest, Some (mk (Ast.Var "false"))))
      else mk (Ast.If (expr x, mk (Ast.Var "true"), Some (go rest)))
  in
  go args

and place (f : Form.t) : Ast.place =
  match f.v with
  | Sym s -> Ast.Pvar s
  | List ({ v = Sym s; _ } :: args)
    when String.length s > 1 && s.[0] = '.' ->
    let field = String.sub s 1 (String.length s - 1) in
    (match args with
     | [ target ] -> Ast.Pfield (expr target, field)
     | _ -> fail f "field place is (.%s value)" field)
  | List ({ v = Sym "at"; _ } :: target :: idx) when idx <> [] ->
    Ast.Pindex (expr target, List.map expr idx)
  | List [ { v = Sym "get"; _ }; m; k ] -> Ast.Pkey (expr m, expr k)
  | List [ { v = Sym "deref"; _ }; p ] -> Ast.Pderef (expr p)
  | _ ->
    fail f
      "%s is not assignable. set takes a name, (.field x), (at a i ...), \
       (get m k) or (deref p)"
      (Form.to_string f)

and arms f (items : Form.t list) : Ast.arm list =
  let rec go = function
    | [] -> []
    | p :: body :: rest ->
      { Ast.pat = pattern p; body = [ expr body ]; aloc = p.loc } :: go rest
    | [ odd ] ->
      Loc.fail odd.loc "match arm %s has no body" (Form.to_string odd)
  in
  if items = [] then Loc.fail f.loc "match needs at least one arm" else go items

and pattern (f : Form.t) : Ast.pattern =
  match f.v with
  | Sym "_" -> Ast.Pwild
  | Kw "else" -> Ast.Pwild
  | Sym ctor -> Ast.Pctor (ctor, [])
  | List ({ v = Sym ctor; _ } :: binds) ->
    Ast.Pctor (ctor, List.map sym binds)
  | _ -> fail f "expected a pattern, found %s" (Form.to_string f)

(* ── Declarations ──────────────────────────────────────────────────── *)

let rec decl types (f : Form.t) : Ast.decl =
  let mk d = { Ast.d; dloc = f.loc } in
  match f.v with
  | List ({ v = Sym "package"; _ } :: args) ->
    (match args with
     | [ n ] -> mk (Ast.Package (sym n))
     | _ -> fail f "package is (package name)")

  | List ({ v = Sym "import"; _ } :: args) ->
    (match args with
     | [ alias; { v = Str path; _ } ] -> mk (Ast.Import (sym alias, path))
     | _ -> fail f "import is (import alias \"collection:path\")")

  | List ({ v = Sym "defalias"; _ } :: args) ->
    (match args with
     | [ n; t ] -> mk (Ast.Defalias (sym n, texpr t))
     | _ -> fail f "defalias is (defalias Name Type)")

  | List ({ v = Sym "defstruct"; _ } :: args) ->
    (match args with
     | [ n; { v = Vec fs; _ } ] -> mk (Ast.Defstruct (sym n, fields f fs))
     | _ -> fail f "defstruct is (defstruct Name [field Type ...])")

  | List ({ v = Sym "defunion"; _ } :: args) ->
    (match args with
     | [ n; { v = Vec vs; _ } ] -> mk (Ast.Defunion (sym n, List.map variant vs))
     | _ -> fail f "defunion is (defunion Name [(Case [field Type ...]) ...])")

  | List ({ v = Sym "defn"; _ } :: args) ->
    (match args with
     | n :: { v = Vec ps; _ } :: rest ->
       let ret, body =
         match rest with
         (* An omitted return type means Unit. A leading form that is a type
            and is not the whole body is the return type. *)
         | [] -> None, []
         | first :: more when more <> [] && is_type_form types first ->
           Some (texpr first), body_of more
         | _ -> None, body_of rest
       in
       mk (Ast.Defn { Ast.name = sym n; params = fields f ps; ret;
                      fbody = body; nloc = n.loc })
     | _ -> fail f "defn is (defn name [param Type ...] ReturnType? body ...)")

  | List ({ v = Sym "declare"; _ } :: args) ->
    (* (declare name [param Type ...] ReturnType? "c_symbol"). The C symbol is
       last and is always written: a foreign name is not derivable from a Flan
       one, and guessing it would fail at link time rather than here. *)
    (match List.rev args with
     | { v = Str csym; _ } :: rest ->
       (match List.rev rest with
        | [ n; { v = Form.Vec ps; _ } ] ->
          mk (Ast.Declare ({ Ast.name = sym n; params = fields f ps;
                             ret = None; fbody = []; nloc = n.loc }, csym))
        | [ n; { v = Form.Vec ps; _ }; r ] ->
          mk (Ast.Declare ({ Ast.name = sym n; params = fields f ps;
                             ret = Some (texpr r); fbody = []; nloc = n.loc },
                           csym))
        | _ -> fail f "declare is (declare name [param Type ...] ReturnType? \"c_symbol\")")
     | _ -> fail f "declare is (declare name [param Type ...] ReturnType? \"c_symbol\")")

  | List ({ v = Sym "defenum"; _ } :: args) ->
    (match args with
     | [ n; { v = Form.Vec ms; _ } ] ->
       let rec pairs = function
         | [] -> []
         | { v = Form.Sym m; _ } :: { v = Form.Int k; _ } :: rest ->
           (m, k) :: pairs rest
         | bad :: _ ->
           fail bad "an enum member is a name followed by an integer, found %s"
             (Form.to_string bad)
       in
       mk (Ast.Defenum (sym n, pairs ms))
     | _ -> fail f "defenum is (defenum Name [member value ...])")

  | List ({ v = Sym "defvar"; _ } :: args) ->
    (match args with
     | [ n; t ] -> mk (Ast.Defvar (sym n, Some (texpr t), Ast.Zeroed))
     | [ n; t; { v = Sym "uninit"; _ } ] ->
       mk (Ast.Defvar (sym n, Some (texpr t), Ast.Uninit))
     | [ n; t; v ] -> mk (Ast.Defvar (sym n, Some (texpr t), Ast.Init (expr v)))
     | _ -> fail f "defvar is (defvar name Type value?)")

  | List ({ v = Sym "defconst"; _ } :: args) ->
    (match args with
     | [ n; v ] -> mk (Ast.Defconst (sym n, None, expr v))
     | [ n; t; v ] -> mk (Ast.Defconst (sym n, Some (texpr t), expr v))
     | _ -> fail f "defconst is (defconst name Type? value)")

  | List ({ v = Sym s; _ } :: _) -> fail f "unknown top-level form (%s ...)" s
  | _ -> fail f "expected a top-level declaration, found %s" (Form.to_string f)

(* Is this form the function's return type, or the first form of its body?

   Shape alone cannot tell: [(Option f64)] and [(Some 1)] are identical
   s-expressions, one a type application and one a constructor call. Capitalised
   heads are not a good enough signal — [(defn f [] (Some 1) (bar))] would eat
   the body's first form as a return type, silently.

   So the decision uses the set of names that are actually types, which the
   pre-pass in [program] collects from the file's own declarations. That makes
   it exact rather than heuristic, because types are only ever introduced by
   defstruct, defunion and defalias — all syntactically obvious. *)
and is_type_form types (f : Form.t) =
  match f.v with
  | Sym s -> Names.mem s types
  | Vec _ -> true                            (* [T] and [n T] are only types *)
  | Map _ -> true                            (* {K V} in this position       *)
  | List ({ v = Sym n; _ } :: _) -> Names.mem n types
  | _ -> false

and variant (f : Form.t) : Ast.variant =
  match f.v with
  | Sym n -> { Ast.vname = n; vfields = []; vloc = f.loc }
  | List [ { v = Sym n; _ }; { v = Vec fs; _ } ] ->
    { Ast.vname = n; vfields = fields f fs; vloc = f.loc }
  | List [ { v = Sym n; _ } ] -> { Ast.vname = n; vfields = []; vloc = f.loc }
  | _ -> fail f "a union case is Name or (Name [field Type ...])"

(* Names introduced as types by this file, plus the builtins. Collected before
   anything is parsed, so a type declared at the bottom of a file is still known
   to a function at the top — top-level names are order-independent. *)
let declared_types (forms : Form.t list) : Names.t =
  List.fold_left
    (fun acc (f : Form.t) ->
       match f.v with
       | List [ { v = Sym ("defstruct" | "defunion" | "defalias"); _ };
                { v = Sym n; _ }; _ ] -> Names.add n acc
       | _ -> acc)
    builtin_types forms

let program (forms : Form.t list) : Ast.decl list =
  let types = declared_types forms in
  List.map (decl types) forms

(* Single-declaration entry point, for tests and the REPL. Sees only the
   builtin types plus whatever this one form declares. *)
let decl (f : Form.t) : Ast.decl = decl (declared_types [ f ]) f