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

(** Why3 / WhyML emitter (MVP, formal verification target).

    Lowers Int/Bool functions to WhyML inside a single theory module.
    Output is consumable by [why3 prove] / [why3 ide] for proof
    obligation generation against Z3 / Alt-Ergo / CVC. *)

open Ast

let why3_reserved = [
  "abstract"; "absurd"; "alias"; "any"; "as"; "assert"; "assume"; "at";
  "axiom"; "begin"; "break"; "by"; "check"; "clone"; "coinductive";
  "constant"; "continue"; "diverges"; "do"; "done"; "downto"; "else";
  "end"; "ensures"; "epsilon"; "exception"; "exists"; "export"; "false";
  "float"; "for"; "forall"; "fun"; "function"; "ghost"; "goal"; "if";
  "import"; "in"; "inductive"; "invariant"; "label"; "lemma"; "let";
  "match"; "meta"; "module"; "mutable"; "not"; "old"; "predicate"; "private";
  "pure"; "raise"; "raises"; "range"; "reads"; "rec"; "ref"; "requires";
  "return"; "returns"; "scope"; "so"; "then"; "theory"; "to"; "true"; "try";
  "type"; "use"; "val"; "variant"; "while"; "with"; "writes";
]

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

let rec why3_type = function
  | TyCon id when id.name = "Int"   -> "int"
  | TyCon id when id.name = "Bool"  -> "bool"
  | TyCon id when id.name = "Float" -> "real"
  | TyCon id when id.name = "Unit"  -> "unit"
  | TyCon id -> mangle id.name
  | TyOwn t | TyRef (_, t) | TyMut (_, t) -> why3_type t
  | _ -> "int"

let ret_type = function None -> "unit" | Some t -> why3_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) -> "(not " ^ 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 in %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 in %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 (see lean_codegen). *)
  if fn_body_contains_return fd.fd_body then
    failwith
      (Printf.sprintf
         "Why3 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) (why3_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 "  let %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 (Why3 / WhyML) *)\n";
  Buffer.add_string buf "(* SPDX-License-Identifier: MPL-2.0 *)\n\n";
  Buffer.add_string buf "module Generated\n";
  Buffer.add_string buf "  use int.Int\n";
  Buffer.add_string buf "  use bool.Bool\n";
  Buffer.add_string buf "  use real.Real\n\n";
  List.iter (function
    | TopFn fd -> Buffer.add_string buf (gen_function fd)
    | _ -> ()
  ) program.prog_decls;
  Buffer.add_string buf "end\n";
  Buffer.contents buf

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