#!/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.
    Strips :effects [...] annotations from defines."""
    from shared.sx.types import Keyword
    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])
                # Strip :effects [...] annotation if present
                # (define name :effects [...] body) → (define name body)
                cleaned = list(expr)
                if (len(cleaned) >= 4 and isinstance(cleaned[2], Keyword)
                        and cleaned[2].name == "effects"):
                    cleaned = [cleaned[0], cleaned[1]] + cleaned[4:]
                defines.append((name, cleaned))
    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 + reactive subscriber fix
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

(* Strict mode refs — used by test runner, stubbed here *)
let _strict_ref = ref Nil
let _prim_param_types_ref = ref Nil
let value_matches_type_p _v _t = Bool true

(* Override reactive_shift_deref to wrap subscriber as NativeFn.
   The transpiler emits bare OCaml closures for (fn () ...) but
   signal_add_sub_b expects SX values. *)
let reactive_shift_deref sig' env kont =
  let scan_result = kont_capture_to_reactive_reset kont in
  let captured_frames = first scan_result in
  let reset_frame = nth scan_result (Number 1.0) in
  let remaining_kont = nth scan_result (Number 2.0) in
  let update_fn = get reset_frame (String "update-fn") in
  let sub_disposers = ref (List []) in
  let subscriber_fn () =
    List.iter (fun d -> ignore (cek_call d Nil)) (sx_to_list !sub_disposers);
    sub_disposers := List [];
    let new_reset = make_reactive_reset_frame env update_fn (Bool false) in
    let new_kont = prim_call "concat" [captured_frames; List [new_reset]; remaining_kont] in
    ignore (with_island_scope
      (fun d -> sub_disposers := sx_append_b !sub_disposers d; Nil)
      (fun () -> cek_run (make_cek_value (signal_value sig') env new_kont)));
    Nil
  in
  let subscriber = NativeFn ("reactive-subscriber", fun _args -> subscriber_fn ()) in
  ignore (signal_add_sub_b sig' subscriber);
  ignore (register_in_scope (fun () ->
    ignore (signal_remove_sub_b sig' subscriber);
    List.iter (fun d -> ignore (cek_call d Nil)) (sx_to_list !sub_disposers);
    Nil));
  let initial_kont = prim_call "concat" [captured_frames; List [reset_frame]; remaining_kont] in
  make_cek_value (signal_value sig') env initial_kont

"""


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/fixups or that belong in web module
        skip = {"trampoline",
                # Freeze functions depend on signals.sx (web spec)
                "freeze-registry", "freeze-signal", "freeze-scope",
                "cek-freeze-scope", "cek-freeze-all",
                "cek-thaw-scope", "cek-thaw-all",
                "freeze-to-sx", "thaw-from-sx",
                "freeze-to-cid", "thaw-from-cid",
                "content-hash", "content-put", "content-get", "content-store"}
        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)


WEB_PREAMBLE = """\
(* sx_web.ml — Auto-generated from web adapters by hosts/ocaml/bootstrap.py *)
(* Do not edit — regenerate with: python3 hosts/ocaml/bootstrap.py --web *)

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

open Sx_types
open Sx_runtime

"""

# Web adapter files to transpile (dependency order)
WEB_ADAPTER_FILES = [
    ("signals.sx", "signals (reactive signal runtime)"),
    ("deps.sx", "deps (component dependency analysis)"),
    ("page-helpers.sx", "page-helpers (pure data transformation helpers)"),
    ("router.sx", "router (client-side route matching)"),
    ("adapter-html.sx", "adapter-html (HTML rendering adapter)"),
]


def compile_web_to_ml(web_dir: str | None = None) -> str:
    """Compile web adapter SX files to OCaml source."""
    from shared.sx.ref.sx_ref import eval_expr, trampoline, make_env, sx_parse

    if web_dir is None:
        web_dir = os.path.join(_PROJECT, "web")

    # 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))

    # Also load the evaluator defines so the transpiler knows about them
    spec_dir = os.path.join(_PROJECT, "spec")
    eval_path = os.path.join(spec_dir, "evaluator.sx")
    if os.path.exists(eval_path):
        with open(eval_path) as f:
            eval_defines = extract_defines(f.read())
        eval_names = [n for n, _ in eval_defines]
    else:
        eval_names = []

    parts = [WEB_PREAMBLE]

    # Collect all web adapter defines
    all_defines = []
    for filename, label in WEB_ADAPTER_FILES:
        filepath = os.path.join(web_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)

        # Deduplicate within file
        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]

        all_defines.extend(defines)
        print(f"  {filename}: {len(defines)} defines", file=sys.stderr)

    # Deduplicate across files (last wins)
    seen = {}
    for i, (n, e) in enumerate(all_defines):
        seen[n] = i
    all_defines = [(n, e) for i, (n, e) in enumerate(all_defines) if seen[n] == i]

    print(f"  Total: {len(all_defines)} unique defines", file=sys.stderr)

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

    # Known defines = evaluator names + web adapter names
    env["_known_defines"] = eval_names + [name for name, _ in all_defines]

    # Translate
    translate_expr = sx_parse("(ml-translate-file _defines)")[0]
    result = trampoline(eval_expr(translate_expr, env))

    parts.append("\n(* === Transpiled from web adapters === *)\n")
    parts.append(result)

    # Registration function — extract actual OCaml names from transpiled output
    # by using the same transpiler mangling.
    # Ask the transpiler for the mangled name of each define.
    name_map = {}
    for name, _ in all_defines:
        mangle_expr = sx_parse(f'(ml-mangle "{name}")')[0]
        mangled = trampoline(eval_expr(mangle_expr, env))
        name_map[name] = mangled

    def count_params(expr):
        """Count actual params from a (define name [annotations] (fn (params...) body)) form."""
        # Find the (fn ...) form — it might be at index 2, 3, or 4 depending on annotations
        fn_expr = None
        for i in range(2, min(len(expr), 6)):
            if (isinstance(expr[i], list) and expr[i] and
                isinstance(expr[i][0], Symbol) and expr[i][0].name in ("fn", "lambda")):
                fn_expr = expr[i]
                break
        if fn_expr is None:
            return -1  # not a function
        params = fn_expr[1] if isinstance(fn_expr[1], list) else []
        n = 0
        skip = False
        for p in params:
            if skip:
                skip = False
                continue
            if isinstance(p, Symbol) and p.name in ("&key", "&rest"):
                skip = True
                continue
            if isinstance(p, list) and len(p) >= 3:  # (name :as type)
                n += 1
            elif isinstance(p, Symbol):
                n += 1
        return n

    parts.append("\n\n(* Register all web adapter functions into an environment *)\n")
    parts.append("let register_web_adapters env =\n")
    for name, expr in all_defines:
        mangled = name_map[name]
        n = count_params(expr)
        if n < 0:
            # Non-function define (constant)
            parts.append(f'  ignore (Sx_types.env_bind env "{name}" {mangled});\n')
        elif n == 0:
            parts.append(f'  ignore (Sx_types.env_bind env "{name}" '
                         f'(NativeFn ("{name}", fun _args -> {mangled} Nil)));\n')
        else:
            # Generate match with correct arity
            arg_names = [chr(97 + i) for i in range(n)]  # a, b, c, ...
            pat = "; ".join(arg_names)
            call = " ".join(arg_names)
            # Pad with Nil for partial application
            pad_call = " ".join(arg_names[:1] + ["Nil"] * (n - 1)) if n > 1 else arg_names[0]
            parts.append(f'  ignore (Sx_types.env_bind env "{name}" '
                         f'(NativeFn ("{name}", fun args -> match args with '
                         f'| [{pat}] -> {mangled} {call} '
                         f'| _ -> raise (Eval_error "{name}: expected {n} args"))));\n')
    parts.append("  ()\n")

    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)",
    )
    parser.add_argument(
        "--web",
        action="store_true",
        help="Compile web adapters instead of evaluator spec",
    )
    parser.add_argument(
        "--web-output",
        default=None,
        help="Output file for web adapters (default: stdout)",
    )
    args = parser.parse_args()

    if args.web or args.web_output:
        result = compile_web_to_ml()
        out = args.web_output or args.output
        if out:
            with open(out, "w") as f:
                f.write(result)
            size = os.path.getsize(out)
            print(f"Wrote {out} ({size} bytes)", file=sys.stderr)
        else:
            print(result)
    else:
        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()