(* SPDX-License-Identifier: MPL-2.0 *)
(* SPDX-FileCopyrightText: 2026 Jonathan D.A. Jewell *)

(** Lean 4 emitter (MVP, dependent-type proof target).

    Lowers Int/Bool/Float functions to Lean 4 [def] declarations. *)

open Ast

let lean_reserved = [
  "abbrev"; "axiom"; "by"; "calc"; "class"; "def"; "deriving"; "do"; "elab";
  "else"; "end"; "example"; "fun"; "have"; "if"; "import"; "in"; "inductive";
  "instance"; "let"; "macro"; "match"; "mut"; "namespace"; "open"; "section";
  "set_option"; "show"; "structure"; "syntax"; "then"; "theorem"; "where";
  "with"; "true"; "false";
]

let mangle s = if List.mem s lean_reserved then s ^ "_" else s

let rec lean_type = function
  | TyCon id when id.name = "Int"   -> "Int"
  | TyCon id when id.name = "Bool"  -> "Bool"
  | TyCon id when id.name = "Float" -> "Float"
  | TyCon id when id.name = "String" -> "String"
  | TyCon id when id.name = "Unit"  -> "Unit"
  | TyCon id -> mangle id.name
  | TyOwn t | TyRef (_, t) | TyMut (_, t) -> lean_type t
  | _ -> "Int"

let ret_type = function None -> "Unit" | Some t -> lean_type t

let rec gen_expr (e : expr) : string =
  match e with
  | ExprLit lit -> gen_lit lit
  | ExprVar id -> mangle id.name
  | ExprApp (callee, args) ->
      let f = gen_expr callee in
      let xs = List.map (fun a -> "(" ^ gen_expr a ^ ")") args in
      f ^ " " ^ String.concat " " xs
  | ExprBinary (a, op, b) ->
      let s = match op with
        | OpAdd -> "+" | OpSub -> "-" | OpMul -> "*" | OpDiv -> "/" | OpMod -> "%"
        | OpEq  -> "==" | OpNe -> "!="
        | OpLt  -> "<" | OpLe -> "<=" | OpGt -> ">" | OpGe -> ">="
        | OpAnd -> "&&" | OpOr -> "||"
        | OpConcat -> "++"
        | _ -> "+"
      in
      "(" ^ gen_expr a ^ " " ^ s ^ " " ^ gen_expr b ^ ")"
  | ExprUnary (OpNeg, x) -> "(- " ^ gen_expr x ^ ")"
  | ExprUnary (OpNot, x) -> "(! " ^ gen_expr x ^ ")"
  | ExprUnary _ -> "0"
  | ExprIf { ei_cond; ei_then; ei_else } ->
      let f = match ei_else with Some e -> gen_expr e | None -> "()" in
      Printf.sprintf "(if %s then %s else %s)" (gen_expr ei_cond) (gen_expr ei_then) f
  | ExprLet { el_pat; el_value; el_body; _ } ->
      let var = match el_pat with PatVar id -> mangle id.name | _ -> "_" in
      let body = match el_body with Some e -> gen_expr e | None -> "()" in
      Printf.sprintf "(let %s := %s; %s)" var (gen_expr el_value) body
  | ExprBlock blk -> gen_block blk
  | ExprSpan (inner, _) -> gen_expr inner
  | ExprReturn (Some e) -> gen_expr e
  | _ -> "0"

and gen_lit = function
  | LitInt (n, _) -> string_of_int n
  | LitFloat (f, _) ->
      let s = string_of_float f in
      if String.length s > 0 && s.[String.length s - 1] = '.' then s ^ "0" else s
  | LitBool (true, _) -> "true"
  | LitBool (false, _) -> "false"
  | LitString (s, _) -> "\"" ^ String.escaped s ^ "\""
  | LitChar (c, _) -> "'" ^ Char.escaped c ^ "'"
  | LitUnit _ -> "()"

and gen_block (blk : block) : string =
  let rec fold = function
    | [] -> (match blk.blk_expr with Some e -> gen_expr e | None -> "()")
    | StmtLet { sl_pat = PatVar id; sl_value; _ } :: rest ->
        Printf.sprintf "(let %s := %s; %s)" (mangle id.name) (gen_expr sl_value) (fold rest)
    | _ :: rest -> fold rest
  in
  fold blk.blk_stmts

let gen_function (fd : fn_decl) : string =
  let name = mangle fd.fd_name.name in
  (* #624: fail loud on `return` rather than silently drop it. The block emitter
     skips non-`let` statements, so an early `return` would vanish (the body
     would lower to `:= ()`). Refuse instead — these are experimental backends. *)
  if fn_body_contains_return fd.fd_body then
    failwith
      (Printf.sprintf
         "Lean backend does not support `return` (in `%s`, Refs #624): early \
          returns would silently drop control flow; rewrite as a tail expression"
         fd.fd_name.name);
  let params = match fd.fd_params with
    | [] -> ""
    | _  -> " " ^ String.concat " " (List.map (fun (p : param) ->
              Printf.sprintf "(%s : %s)" (mangle p.p_name.name) (lean_type p.p_ty))
              fd.fd_params)
  in
  let ret = ret_type fd.fd_ret_ty in
  let body = match fd.fd_body with
    | FnExpr e -> gen_expr e
    | FnBlock b -> gen_block b
  in
  Printf.sprintf "def %s%s : %s :=\n  %s\n\n" name params ret body

let generate (program : program) (_symbols : Symbol.t) : string =
  let buf = Buffer.create 1024 in
  Buffer.add_string buf "-- Generated by AffineScript compiler (Lean 4)\n";
  Buffer.add_string buf "-- SPDX-License-Identifier: MPL-2.0\n\n";
  List.iter (function
    | TopFn fd -> Buffer.add_string buf (gen_function fd)
    | _ -> ()
  ) program.prog_decls;
  Buffer.contents buf

let codegen_lean (program : program) (symbols : Symbol.t) : (string, string) result =
  try Ok (generate program symbols)
  with
  | Failure m -> Error ("Lean codegen error: " ^ m)
  | e         -> Error ("Lean codegen error: " ^ Printexc.to_string e)
