#!/usr/bin/env python3
"""
Bootstrap compiler: SX spec -> OCaml.

Loads the SX-to-OCaml transpiler (transpiler.sx), feeds it the spec files,
and produces sx_ref.ml — the transpiled evaluator as native OCaml.

Usage:
    python3 hosts/ocaml/bootstrap.py --output hosts/ocaml/lib/sx_ref.ml
"""
from __future__ import annotations

import os
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
_PROJECT = os.path.abspath(os.path.join(_HERE, "..", ".."))
sys.path.insert(0, _PROJECT)

from shared.sx.parser import parse_all
from shared.sx.types import Symbol


def extract_defines(source: str) -> list[tuple[str, list]]:
    """Parse .sx source, return list of (name, define-expr) for top-level defines."""
    exprs = parse_all(source)
    defines = []
    for expr in exprs:
        if isinstance(expr, list) and expr and isinstance(expr[0], Symbol):
            if expr[0].name == "define":
                name = expr[1].name if isinstance(expr[1], Symbol) else str(expr[1])
                defines.append((name, expr))
    return defines


# OCaml preamble — opens and runtime helpers
PREAMBLE = """\
(* sx_ref.ml — Auto-generated from SX spec by hosts/ocaml/bootstrap.py *)
(* Do not edit — regenerate with: python3 hosts/ocaml/bootstrap.py *)

[@@@warning "-26-27"]

open Sx_types
open Sx_runtime

(* Trampoline — evaluates thunks via the CEK machine.
   eval_expr is defined in the transpiled block below. *)
let trampoline v = v  (* CEK machine doesn't produce thunks *)

"""


# OCaml fixups — override iterative CEK run
FIXUPS = """\

(* Override recursive cek_run with iterative loop *)
let cek_run_iterative state =
  let s = ref state in
  while not (match cek_terminal_p !s with Bool true -> true | _ -> false) do
    s := cek_step !s
  done;
  cek_value !s

"""


def compile_spec_to_ml(spec_dir: str | None = None) -> str:
    """Compile the SX spec to OCaml source."""
    from shared.sx.ref.sx_ref import eval_expr, trampoline, make_env, sx_parse

    if spec_dir is None:
        spec_dir = os.path.join(_PROJECT, "spec")

    # Load the transpiler
    env = make_env()
    transpiler_path = os.path.join(_HERE, "transpiler.sx")
    with open(transpiler_path) as f:
        transpiler_src = f.read()
    for expr in sx_parse(transpiler_src):
        trampoline(eval_expr(expr, env))

    # Spec files to transpile (in dependency order)
    sx_files = [
        ("evaluator.sx", "evaluator (frames + eval + CEK)"),
    ]

    parts = [PREAMBLE]

    for filename, label in sx_files:
        filepath = os.path.join(spec_dir, filename)
        if not os.path.exists(filepath):
            print(f"Warning: {filepath} not found, skipping", file=sys.stderr)
            continue

        with open(filepath) as f:
            src = f.read()
        defines = extract_defines(src)

        # Skip defines provided by preamble or fixups
        skip = {"trampoline"}
        defines = [(n, e) for n, e in defines if n not in skip]

        # Deduplicate — keep last definition for each name (CEK overrides tree-walk)
        seen = {}
        for i, (n, e) in enumerate(defines):
            seen[n] = i
        defines = [(n, e) for i, (n, e) in enumerate(defines) if seen[n] == i]

        # Build the defines list for the transpiler
        defines_list = [[name, expr] for name, expr in defines]
        env["_defines"] = defines_list

        # Pass known define names so the transpiler can distinguish
        # static (OCaml fn) calls from dynamic (SX value) calls
        env["_known_defines"] = [name for name, _ in defines]

        # Call ml-translate-file — emits as single let rec block
        translate_expr = sx_parse("(ml-translate-file _defines)")[0]
        result = trampoline(eval_expr(translate_expr, env))

        parts.append(f"\n(* === Transpiled from {label} === *)\n")
        parts.append(result)

    parts.append(FIXUPS)
    return "\n".join(parts)


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Bootstrap SX spec -> OCaml")
    parser.add_argument(
        "--output", "-o",
        default=None,
        help="Output file (default: stdout)",
    )
    args = parser.parse_args()

    result = compile_spec_to_ml()

    if args.output:
        with open(args.output, "w") as f:
            f.write(result)
        size = os.path.getsize(args.output)
        print(f"Wrote {args.output} ({size} bytes)", file=sys.stderr)
    else:
        print(result)


if __name__ == "__main__":
    main()