#!/usr/bin/env python3
"""
Run SX spec tests using the bootstrapped Python evaluator.

Usage:
    python3 hosts/python/tests/run_tests.py                    # all spec tests
    python3 hosts/python/tests/run_tests.py test-primitives     # specific test
    python3 hosts/python/tests/run_tests.py --full              # include optional modules
"""
from __future__ import annotations
import os, sys

# Increase recursion limit for TCO tests (Python's default 1000 is too low)
sys.setrecursionlimit(5000)

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

from shared.sx.ref.sx_ref import sx_parse as parse_all
from shared.sx.ref import sx_ref
from shared.sx.ref.sx_ref import (
    make_env, env_get, env_has, env_set, env_extend, env_merge,
)
from shared.sx.types import (
    NIL, Symbol, Keyword, Lambda, Component, Island, Macro,
)

# Use tree-walk evaluator
eval_expr = sx_ref._tree_walk_eval_expr
trampoline = sx_ref._tree_walk_trampoline
sx_ref.eval_expr = eval_expr
sx_ref.trampoline = trampoline

# Check for --full flag
full_build = "--full" in sys.argv

# Build env with primitives
env = make_env()

# ---------------------------------------------------------------------------
# Test infrastructure
# ---------------------------------------------------------------------------
_suite_stack: list[str] = []
_pass_count = 0
_fail_count = 0


def _try_call(thunk):
    try:
        trampoline(eval_expr([thunk], env))
        return {"ok": True}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def _report_pass(name):
    global _pass_count
    _pass_count += 1
    ctx = " > ".join(_suite_stack)
    print(f"  PASS: {ctx} > {name}")
    return NIL


def _report_fail(name, error):
    global _fail_count
    _fail_count += 1
    ctx = " > ".join(_suite_stack)
    print(f"  FAIL: {ctx} > {name}: {error}")
    return NIL


def _push_suite(name):
    _suite_stack.append(name)
    print(f"{'  ' * (len(_suite_stack)-1)}Suite: {name}")
    return NIL


def _pop_suite():
    if _suite_stack:
        _suite_stack.pop()
    return NIL


env["try-call"] = _try_call
env["report-pass"] = _report_pass
env["report-fail"] = _report_fail
env["push-suite"] = _push_suite
env["pop-suite"] = _pop_suite

# ---------------------------------------------------------------------------
# Test helpers
# ---------------------------------------------------------------------------


def _deep_equal(a, b):
    if a is b:
        return True
    if a is NIL and b is NIL:
        return True
    if a is NIL or b is NIL:
        return a is None and b is NIL or b is None and a is NIL
    if type(a) != type(b):
        # number comparison: int vs float
        if isinstance(a, (int, float)) and isinstance(b, (int, float)):
            return a == b
        return False
    if isinstance(a, list):
        if len(a) != len(b):
            return False
        return all(_deep_equal(x, y) for x, y in zip(a, b))
    if isinstance(a, dict):
        ka = {k for k in a if k != "_nil"}
        kb = {k for k in b if k != "_nil"}
        if ka != kb:
            return False
        return all(_deep_equal(a[k], b[k]) for k in ka)
    return a == b


env["equal?"] = _deep_equal
env["identical?"] = lambda a, b: a is b


def _test_env():
    return make_env()


def _sx_parse(source):
    return parse_all(source)


def _sx_parse_one(source):
    exprs = parse_all(source)
    return exprs[0] if exprs else NIL


env["test-env"] = _test_env
env["sx-parse"] = _sx_parse
env["sx-parse-one"] = _sx_parse_one
env["cek-eval"] = lambda s: trampoline(eval_expr(parse_all(s)[0], make_env())) if parse_all(s) else NIL
env["eval-expr-cek"] = lambda expr, e=None: trampoline(eval_expr(expr, e or env))

# Env operations
env["env-get"] = env_get
env["env-has?"] = env_has
env["env-set!"] = env_set
env["env-bind!"] = lambda e, k, v: e.__setitem__(k, v) or v
env["env-extend"] = env_extend
env["env-merge"] = env_merge

# Missing primitives
env["upcase"] = lambda s: str(s).upper()
env["downcase"] = lambda s: str(s).lower()
env["make-keyword"] = lambda name: Keyword(name)
env["make-symbol"] = lambda name: Symbol(name)
env["string-length"] = lambda s: len(str(s))
env["dict-get"] = lambda d, k: d.get(k, NIL) if isinstance(d, dict) else NIL
env["apply"] = lambda f, *args: f(*args[-1]) if args and isinstance(args[-1], list) else f()

# Render helpers
def _render_html(src, e=None):
    if isinstance(src, str):
        parsed = parse_all(src)
        if not parsed:
            return ""
        expr = parsed[0] if len(parsed) == 1 else [Symbol("do")] + parsed
        result = sx_ref.render_to_html(expr, e or make_env())
        # Reset render mode
        sx_ref._render_mode = False
        return result
    result = sx_ref.render_to_html(src, e or env)
    sx_ref._render_mode = False
    return result


env["render-html"] = _render_html
env["render-to-html"] = _render_html
env["string-contains?"] = lambda s, sub: str(sub) in str(s)

# Type system helpers
env["test-prim-types"] = lambda: {
    "+": "number", "-": "number", "*": "number", "/": "number",
    "mod": "number", "inc": "number", "dec": "number",
    "abs": "number", "min": "number", "max": "number",
    "str": "string", "upper": "string", "lower": "string",
    "trim": "string", "join": "string", "replace": "string",
    "=": "boolean", "<": "boolean", ">": "boolean",
    "<=": "boolean", ">=": "boolean",
    "not": "boolean", "nil?": "boolean", "empty?": "boolean",
    "number?": "boolean", "string?": "boolean", "boolean?": "boolean",
    "list?": "boolean", "dict?": "boolean",
    "contains?": "boolean", "has-key?": "boolean",
    "starts-with?": "boolean", "ends-with?": "boolean",
    "len": "number", "first": "any", "rest": "list",
    "last": "any", "nth": "any", "cons": "list",
    "append": "list", "concat": "list", "reverse": "list",
    "sort": "list", "slice": "list", "range": "list",
    "flatten": "list", "keys": "list", "vals": "list",
    "assoc": "dict", "dissoc": "dict", "merge": "dict", "dict": "dict",
    "get": "any", "type-of": "string",
}
env["test-prim-param-types"] = lambda: {
    "+": {"positional": [["a", "number"]], "rest-type": "number"},
    "-": {"positional": [["a", "number"]], "rest-type": "number"},
    "*": {"positional": [["a", "number"]], "rest-type": "number"},
    "/": {"positional": [["a", "number"]], "rest-type": "number"},
    "inc": {"positional": [["n", "number"]], "rest-type": NIL},
    "dec": {"positional": [["n", "number"]], "rest-type": NIL},
    "upper": {"positional": [["s", "string"]], "rest-type": NIL},
    "lower": {"positional": [["s", "string"]], "rest-type": NIL},
    "keys": {"positional": [["d", "dict"]], "rest-type": NIL},
    "vals": {"positional": [["d", "dict"]], "rest-type": NIL},
}
env["component-param-types"] = lambda c: getattr(c, "_param_types", NIL)
env["component-set-param-types!"] = lambda c, t: setattr(c, "_param_types", t) or NIL
env["component-params"] = lambda c: c.params
env["component-body"] = lambda c: c.body
env["component-has-children"] = lambda c: c.has_children
env["component-affinity"] = lambda c: getattr(c, "affinity", "auto")

# Type accessors
env["callable?"] = lambda x: callable(x) or isinstance(x, (Lambda, Component, Island))
env["lambda?"] = lambda x: isinstance(x, Lambda)
env["component?"] = lambda x: isinstance(x, Component)
env["island?"] = lambda x: isinstance(x, Island)
env["macro?"] = lambda x: isinstance(x, Macro)
env["thunk?"] = sx_ref.is_thunk
env["thunk-expr"] = sx_ref.thunk_expr
env["thunk-env"] = sx_ref.thunk_env
env["make-thunk"] = sx_ref.make_thunk
env["make-lambda"] = sx_ref.make_lambda
env["make-component"] = sx_ref.make_component
env["make-macro"] = sx_ref.make_macro
env["lambda-params"] = lambda f: f.params
env["lambda-body"] = lambda f: f.body
env["lambda-closure"] = lambda f: f.closure
env["lambda-name"] = lambda f: f.name
env["set-lambda-name!"] = lambda f, n: setattr(f, "name", n) or NIL
env["component-closure"] = lambda c: c.closure
env["component-name"] = lambda c: c.name
env["component-has-children?"] = lambda c: c.has_children
env["macro-params"] = lambda m: m.params
env["macro-rest-param"] = lambda m: m.rest_param
env["macro-body"] = lambda m: m.body
env["macro-closure"] = lambda m: m.closure
env["symbol-name"] = lambda s: s.name if isinstance(s, Symbol) else str(s)
env["keyword-name"] = lambda k: k.name if isinstance(k, Keyword) else str(k)
env["sx-serialize"] = sx_ref.sx_serialize if hasattr(sx_ref, "sx_serialize") else lambda x: str(x)

# Render dispatch globals — evaluator checks *render-check* and *render-fn*
env["*render-check*"] = NIL
env["*render-fn*"] = NIL

# Custom special forms registry — modules register forms at load time
env["*custom-special-forms*"] = {}
def _register_special_form(name, handler):
    env["*custom-special-forms*"][name] = handler
    return NIL
env["register-special-form!"] = _register_special_form

# is-else-clause? — check if a cond/case test is an else marker
def _is_else_clause(test):
    if isinstance(test, Keyword) and test.name == "else":
        return True
    if isinstance(test, Symbol) and test.name in ("else", ":else"):
        return True
    return False
env["is-else-clause?"] = _is_else_clause

# Strict mode stubs (not yet bootstrapped to Python — no-ops for now)
env["set-strict!"] = lambda val: NIL
env["set-prim-param-types!"] = lambda types: NIL
env["value-matches-type?"] = lambda val, t: True
env["*strict*"] = False
env["primitive?"] = lambda name: name in env
env["get-primitive"] = lambda name: env.get(name, NIL)

# ---------------------------------------------------------------------------
# Load test framework
# ---------------------------------------------------------------------------
framework_src = open(os.path.join(_SPEC_TESTS, "test-framework.sx")).read()
for expr in parse_all(framework_src):
    trampoline(eval_expr(expr, env))

# ---------------------------------------------------------------------------
# Determine which tests to run
# ---------------------------------------------------------------------------
args = [a for a in sys.argv[1:] if not a.startswith("--")]

# Tests requiring optional modules (only with --full)
REQUIRES_FULL = {"test-continuations.sx", "test-continuations-advanced.sx", "test-types.sx", "test-freeze.sx", "test-strict.sx", "test-cek.sx", "test-cek-advanced.sx", "test-signals-advanced.sx"}

test_files = []
if args:
    for arg in args:
        name = arg if arg.endswith(".sx") else f"{arg}.sx"
        p = os.path.join(_SPEC_TESTS, name)
        if os.path.exists(p):
            test_files.append(p)
        else:
            print(f"Test file not found: {name}")
else:
    for f in sorted(os.listdir(_SPEC_TESTS)):
        if f.startswith("test-") and f.endswith(".sx") and f != "test-framework.sx":
            if not full_build and f in REQUIRES_FULL:
                print(f"Skipping {f} (requires --full)")
                continue
            test_files.append(os.path.join(_SPEC_TESTS, f))

# ---------------------------------------------------------------------------
# Run tests
# ---------------------------------------------------------------------------
for test_file in test_files:
    name = os.path.basename(test_file)
    print("=" * 60)
    print(f"Running {name}")
    print("=" * 60)
    try:
        src = open(test_file).read()
        exprs = parse_all(src)
        for expr in exprs:
            trampoline(eval_expr(expr, env))
    except Exception as e:
        print(f"ERROR in {name}: {e}")
        _fail_count += 1

# Summary
print("=" * 60)
print(f"Results: {_pass_count} passed, {_fail_count} failed")
print("=" * 60)
sys.exit(1 if _fail_count > 0 else 0)