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

(** Lua 5.x Emitter (MVP).

    Lowers a subset of AffineScript to Lua source. Numbers are unified as
    Lua's native [number] type — Int and Float both map to it. *)

open Ast

let lua_reserved = [
  "and"; "break"; "do"; "else"; "elseif"; "end"; "false"; "for";
  "function"; "goto"; "if"; "in"; "local"; "nil"; "not"; "or"; "repeat";
  "return"; "then"; "true"; "until"; "while";
]

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

let prelude = {|-- AffineScript Lua runtime (MVP)
local function println(s) print(tostring(s)) end
local function as_print(s) io.write(tostring(s)) end
local Some = function(v) return { tag = "Some", value = v } end
local None = { tag = "None" }
local Ok   = function(v) return { tag = "Ok",  value = v } end
local Err  = function(e) return { tag = "Err", error = e } end
local function read_line() local l = io.read("l"); return l or "" end

|}

let rec gen_expr (e : expr) : string =
  match e with
  | ExprLit lit -> gen_lit lit
  | ExprVar id -> mangle id.name
  | ExprApp (callee, args) ->
      gen_expr callee ^ "(" ^ String.concat ", " (List.map gen_expr args) ^ ")"
  | ExprBinary (a, op, b) ->
      let s = match op with
        | OpAdd -> "+" | OpSub -> "-" | OpMul -> "*" | OpDiv -> "/" | OpMod -> "%"
        | OpEq  -> "==" | OpNe -> "~="
        | OpLt  -> "<" | OpLe -> "<=" | OpGt -> ">" | OpGe -> ">="
        | OpAnd -> "and" | OpOr -> "or"
        | OpBitAnd -> "&" | OpBitOr -> "|" | OpBitXor -> "~"  (* Lua 5.3+ *)
        | OpShl -> "<<" | OpShr -> ">>"
        | OpConcat -> ".."
      in
      "(" ^ gen_expr a ^ " " ^ s ^ " " ^ gen_expr b ^ ")"
  | ExprUnary (op, x) ->
      (match op with
       | OpNeg -> "(-" ^ gen_expr x ^ ")"
       | OpNot -> "(not " ^ gen_expr x ^ ")"
       | OpBitNot -> "(~" ^ gen_expr x ^ ")"
       | OpRef -> gen_expr x
       | OpDeref -> gen_expr x)
  | ExprIf { ei_cond; ei_then; ei_else } ->
      (* Lua has no expression-form if; emit a do-block-as-expression via
         IIFE. *)
      let c = gen_expr ei_cond in
      let t = gen_expr ei_then in
      let f = match ei_else with Some e -> gen_expr e | None -> "nil" in
      Printf.sprintf "((function() if %s then return %s else return %s end end)())" c t f
  | ExprLet { el_pat; el_value; el_body; _ } ->
      let var = match el_pat with PatVar id -> mangle id.name | _ -> "_" in
      let v = gen_expr el_value in
      let body = match el_body with Some e -> gen_expr e | None -> "nil" in
      Printf.sprintf "((function() local %s = %s; return %s end)())" var v body
  | ExprBlock blk -> gen_block_expr blk
  | ExprTuple es -> "{" ^ String.concat ", " (List.map gen_expr es) ^ "}"
  | ExprTupleIndex (e, n) ->
      (* Lua tables are 1-indexed; AS tuple indices are 0-based. *)
      Printf.sprintf "%s[%d]" (gen_expr e) (n + 1)
  | ExprArray es -> "{" ^ String.concat ", " (List.map gen_expr es) ^ "}"
  | ExprIndex (a, i) -> gen_expr a ^ "[" ^ gen_expr i ^ "]"
  | ExprRecord { er_fields; _ } ->
      let fs = List.map (fun (id, e_opt) ->
        let v = match e_opt with Some e -> gen_expr e | None -> mangle id.name in
        Printf.sprintf "%s = %s" (mangle id.name) v
      ) er_fields in
      "{ " ^ String.concat ", " fs ^ " }"
  | ExprField (record, field) ->
      gen_expr record ^ "." ^ mangle field.name
  | ExprVariant (_ty, ctor) -> mangle ctor.name
  | ExprMatch { em_scrutinee; em_arms } ->
      let scrut_str = gen_expr em_scrutinee in
      let rec arms = function
        | [] -> "error(\"non-exhaustive match\")"
        | arm :: rest ->
            let cond = gen_pattern_test "__scrut" arm.ma_pat in
            let bindings = gen_pattern_bindings "__scrut" arm.ma_pat in
            let body = gen_expr arm.ma_body in
            Printf.sprintf "if %s then %s return %s else %s end"
              cond bindings body (arms rest)
      in
      Printf.sprintf "((function(__scrut) %s end)(%s))" (arms em_arms) scrut_str
  | ExprSpan (inner, _) -> gen_expr inner
  | ExprReturn (Some e) -> gen_expr e
  | _ -> "error(\"Lua backend: unsupported expression\")"

and gen_pattern_test scrut pat =
  match pat with
  | PatWildcard _ | PatVar _ -> "true"
  | PatLit lit -> Printf.sprintf "%s == %s" scrut (gen_lit lit)
  | PatCon (id, _) -> Printf.sprintf "%s.tag == %S" scrut id.name
  | PatTuple _ -> "true"  (* arity match by structure, not tag *)
  | PatRecord _ -> "true"
  | PatAs (_, p) -> gen_pattern_test scrut p
  | PatOr (p, _) -> gen_pattern_test scrut p

and gen_pattern_bindings scrut pat =
  let buf = Buffer.create 64 in
  let rec walk path = function
    | PatWildcard _ | PatLit _ -> ()
    | PatVar id ->
        Buffer.add_string buf
          (Printf.sprintf "local %s = %s; " (mangle id.name) path)
    | PatTuple ps ->
        List.iteri (fun i p -> walk (Printf.sprintf "%s[%d]" path (i + 1)) p) ps
    | PatRecord (fields, _) ->
        List.iter (fun (id, sub) ->
          let p = path ^ "." ^ mangle id.name in
          match sub with
          | None -> Buffer.add_string buf (Printf.sprintf "local %s = %s; " (mangle id.name) p)
          | Some sub -> walk p sub
        ) fields
    | PatCon (_, args) ->
        (match args with
         | [] -> ()
         | [single] -> walk (path ^ ".value") single
         | many ->
             List.iteri (fun i p ->
               walk (Printf.sprintf "%s.values[%d]" path (i + 1)) p
             ) many)
    | PatAs (id, sub) ->
        Buffer.add_string buf (Printf.sprintf "local %s = %s; " (mangle id.name) path);
        walk path sub
    | PatOr (p, _) -> walk path p
  in
  walk scrut pat;
  Buffer.contents buf

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 _ -> "nil"

and gen_block_expr (blk : block) : string =
  let stmts = String.concat " " (List.map gen_stmt blk.blk_stmts) in
  let tail = match blk.blk_expr with
    | Some e -> "return " ^ gen_expr e
    | None   -> "return nil"
  in
  Printf.sprintf "((function() %s %s end)())" stmts tail

and gen_stmt = function
  | StmtLet { sl_pat = PatVar id; sl_value; _ } ->
      Printf.sprintf "local %s = %s;" (mangle id.name) (gen_expr sl_value)
  | StmtLet _ -> ""
  | StmtExpr e -> gen_expr e ^ ";"
  | StmtAssign (lhs, _, rhs) ->
      Printf.sprintf "%s = %s;" (gen_expr lhs) (gen_expr rhs)
  | StmtWhile (cond, body) ->
      Printf.sprintf "while %s do %s end"
        (gen_expr cond)
        (String.concat " " (List.map gen_stmt body.blk_stmts))
  | StmtFor _ -> ""

let gen_function (fd : fn_decl) : string =
  let name = mangle fd.fd_name.name in
  let params = String.concat ", "
    (List.map (fun (p : param) -> mangle p.p_name.name) fd.fd_params) in
  let body = match fd.fd_body with
    | FnExpr e -> "return " ^ gen_expr e
    | FnBlock b ->
        let stmts = String.concat " " (List.map gen_stmt b.blk_stmts) in
        let tail = match b.blk_expr with
          | Some e -> "return " ^ gen_expr e
          | None   -> ""
        in
        stmts ^ " " ^ tail
  in
  Printf.sprintf "function %s(%s)\n  %s\nend\n" name params body

let gen_type_decl (td : type_decl) : string =
  match td.td_body with
  | TyEnum variants ->
      let buf = Buffer.create 64 in
      List.iter (fun (vd : variant_decl) ->
        let arity = List.length vd.vd_fields in
        let name = mangle vd.vd_name.name in
        if arity = 0 then
          Buffer.add_string buf (Printf.sprintf "%s = { tag = %S }\n" name vd.vd_name.name)
        else if arity = 1 then
          Buffer.add_string buf
            (Printf.sprintf "%s = function(v) return { tag = %S, value = v } end\n"
               name vd.vd_name.name)
        else
          let ps = List.init arity (fun i -> "v" ^ string_of_int i) in
          Buffer.add_string buf
            (Printf.sprintf "%s = function(%s) return { tag = %S, values = {%s} } end\n"
               name (String.concat ", " ps) vd.vd_name.name (String.concat ", " ps))
      ) variants;
      Buffer.add_char buf '\n';
      Buffer.contents buf
  | _ -> ""  (* records erased — Lua tables are duck-typed *)

let generate (program : program) (_symbols : Symbol.t) : string =
  let buf = Buffer.create 1024 in
  Buffer.add_string buf "-- Generated by AffineScript compiler\n";
  Buffer.add_string buf "-- SPDX-License-Identifier: MPL-2.0\n";
  Buffer.add_string buf prelude;
  List.iter (function
    | TopType td -> Buffer.add_string buf (gen_type_decl td)
    | _ -> ()
  ) program.prog_decls;
  List.iter (function
    | TopFn fd -> Buffer.add_string buf (gen_function fd); Buffer.add_char buf '\n'
    | _ -> ()
  ) program.prog_decls;
  let has_main = List.exists (function
    | TopFn fd -> fd.fd_name.name = "main"
    | _ -> false) program.prog_decls in
  if has_main then Buffer.add_string buf "print(main())\n";
  Buffer.contents buf

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