#!/usr/bin/env python3
"""
Acceptance KATs for glyphscope-based vault-anchored CDC code generation.

Tests:
- Vocab KAT: All SAT tokens exist in CODE_VOCAB
- Header KAT: Scaffold produces exact function signature
- State KAT: Spot checks on FSM state transitions
- Macro KAT: @LEV_DP expansion SHA256 matches
- Compile KAT: Generated code compiles as valid Python
- Determinism KAT: Same seed produces identical output
"""

import os
import sys
import hashlib

# Allow fallback anchor for testing
os.environ["VAULT_DISABLED"] = "1"

from helix_cdc.glyphscope_code import (
    CODE_VOCAB, SCAFFOLD, BODY_HINT, LEV_DP_BODY, LEV_DP_SHA256,
    id_of, is_letter, ids_to_code, CodeSAT
)
from helix_cdc.cdc_decode import CDCDecoder


def test_vocab_kat():
    """KAT 1: All SAT tokens exist in CODE_VOCAB"""
    print("=== Vocab KAT ===")

    # Check scaffold tokens
    missing = []
    for tok in SCAFFOLD:
        try:
            idx = id_of(tok)
            assert CODE_VOCAB[idx] == tok
        except (ValueError, AssertionError):
            missing.append(tok)

    # Check body hint tokens
    for tok in BODY_HINT:
        try:
            idx = id_of(tok)
            assert CODE_VOCAB[idx] == tok
        except (ValueError, AssertionError):
            if tok not in missing:
                missing.append(tok)

    # Check macro
    try:
        idx = id_of("@LEV_DP")
        assert CODE_VOCAB[idx] == "@LEV_DP"
    except (ValueError, AssertionError):
        missing.append("@LEV_DP")

    if missing:
        print(f"❌ FAIL: Missing tokens in CODE_VOCAB: {missing}")
        return False

    print(f"✅ PASS: All {len(set(SCAFFOLD + BODY_HINT + ['@LEV_DP']))} required tokens exist")
    return True


def test_header_kat():
    """KAT 2: Scaffold produces exact function header"""
    print("\n=== Header KAT ===")

    decoder = CDCDecoder(seed_path=None)
    text, trace = decoder.generate_from_anchor("", max_new_tokens=64, temperature=0.1, seed=123)

    expected_header = "def levenshtein(a, b):\n    "
    if not text.startswith(expected_header):
        print(f"❌ FAIL: Header mismatch")
        print(f"Expected: {expected_header!r}")
        print(f"Got: {text[:len(expected_header)]!r}")
        return False

    scaffold_used = trace["scaffold_tokens_used"]
    if scaffold_used != "22/22":
        print(f"❌ FAIL: Scaffold not fully used: {scaffold_used}")
        return False

    print(f"✅ PASS: Scaffold produced exact header ({scaffold_used} tokens)")
    return True


def test_state_kat():
    """KAT 3: Spot checks on FSM state transitions"""
    print("\n=== State KAT ===")

    sat = CodeSAT(use_scaffold=False)  # Test raw FSM without scaffold

    # Test 1: START → only "def" allowed
    allowed = sat.allowed_ids([])
    if allowed != [id_of("def")]:
        print(f"❌ FAIL: START should only allow 'def', got {len(allowed)} tokens")
        return False

    # Test 2: After "def" → space + letters
    allowed = sat.allowed_ids(["def"])
    if id_of(" ") not in allowed:
        print(f"❌ FAIL: After 'def' should allow space")
        return False
    if id_of("a") not in allowed:
        print(f"❌ FAIL: After 'def' should allow letters")
        return False

    # Test 3: After "(" → letters + space + ")"
    sat.state = "IN_PARAMS"
    allowed = sat.allowed_ids(["("])
    if id_of("a") not in allowed or id_of(")") not in allowed:
        print(f"❌ FAIL: After '(' should allow letters and ')'")
        return False

    # Test 4: After ":" → only newline
    sat.state = "AFTER_COLON"
    allowed = sat.allowed_ids([":"])
    if allowed != [id_of("\n")]:
        print(f"❌ FAIL: After ':' should only allow newline")
        return False

    # Test 5: After indent → return/if/@LEV_DP
    sat.state = "AFTER_NEWLINE"
    allowed = sat.allowed_ids(["    "])
    expected = {id_of("return"), id_of("if"), id_of("@LEV_DP")}
    if set(allowed) != expected:
        print(f"❌ FAIL: After indent should allow return/if/@LEV_DP")
        return False

    print("✅ PASS: All 5 state transitions correct")
    return True


def test_macro_kat():
    """KAT 4: @LEV_DP expansion SHA256 matches"""
    print("\n=== Macro KAT ===")

    # Test expansion
    ids = [id_of("@LEV_DP")]
    expanded = ids_to_code(ids)

    if expanded != LEV_DP_BODY:
        print(f"❌ FAIL: Macro expansion mismatch")
        print(f"Expected length: {len(LEV_DP_BODY)}")
        print(f"Got length: {len(expanded)}")
        return False

    # Test SHA256
    computed_sha = hashlib.sha256(LEV_DP_BODY.encode()).hexdigest()
    if computed_sha != LEV_DP_SHA256:
        print(f"❌ FAIL: SHA256 mismatch")
        print(f"Expected: {LEV_DP_SHA256[:16]}...")
        print(f"Got: {computed_sha[:16]}...")
        return False

    print(f"✅ PASS: Macro expands to {len(LEV_DP_BODY)} bytes, SHA256={LEV_DP_SHA256[:16]}...")
    return True


def test_compile_kat():
    """KAT 5: Generated code compiles as valid Python"""
    print("\n=== Compile KAT ===")

    decoder = CDCDecoder(seed_path=None)
    text, trace = decoder.generate_from_anchor("", max_new_tokens=64, temperature=0.3, seed=42)

    if not trace["compiled"]:
        print(f"❌ FAIL: Code does not compile")
        print(f"Error: {trace.get('compile_error', 'Unknown')}")
        print(f"Code:\n{text}")
        return False

    # Test it actually works
    namespace = {}
    try:
        exec(text, namespace)
        levenshtein = namespace.get("levenshtein")
        if not callable(levenshtein):
            print(f"❌ FAIL: No callable 'levenshtein' function")
            return False

        # Quick sanity test
        result = levenshtein("kitten", "sitting")
        if result != 3:
            print(f"❌ FAIL: levenshtein('kitten', 'sitting') = {result}, expected 3")
            return False
    except Exception as e:
        print(f"❌ FAIL: Runtime error: {e}")
        return False

    print(f"✅ PASS: Code compiles and executes correctly")
    print(f"   Macro used: {trace['macro_used']}")
    print(f"   Tokens: {trace['generated_tokens']} (scaffold + macro)")
    return True


def test_determinism_kat():
    """KAT 6: Same seed produces identical output"""
    print("\n=== Determinism KAT ===")

    decoder = CDCDecoder(seed_path=None)

    # Generate 3 times with same seed
    runs = []
    for i in range(3):
        text, trace = decoder.generate_from_anchor("", max_new_tokens=64, temperature=0.3, seed=999)
        runs.append({
            "text": text,
            "symbols": trace["raw_symbols"],
            "tiles": [step["tile_sha256"] for step in trace["tile_stream"]],
        })

    # Check all match
    if not all(r["text"] == runs[0]["text"] for r in runs):
        print(f"❌ FAIL: Text differs across runs")
        return False

    if not all(r["symbols"] == runs[0]["symbols"] for r in runs):
        print(f"❌ FAIL: Raw symbols differ across runs")
        return False

    if not all(r["tiles"] == runs[0]["tiles"] for r in runs):
        print(f"❌ FAIL: Tile SHA256s differ across runs")
        return False

    print(f"✅ PASS: 3 runs with seed=999 produce identical output")
    print(f"   Text length: {len(runs[0]['text'])} bytes")
    print(f"   Tiles: {len(runs[0]['tiles'])} steps")
    return True


def test_no_macro_kat():
    """KAT 7: No-macro mode (GLYPHSCOPE_ALLOW_MACROS=0) works with body_hint fallback"""
    print("\n=== No-Macro KAT ===")

    # Enable no-macro mode
    os.environ["GLYPHSCOPE_ALLOW_MACROS"] = "0"

    decoder = CDCDecoder(seed_path=None)
    text, trace = decoder.generate_from_anchor("", max_new_tokens=64, temperature=0.3, seed=123)

    # Reset env
    del os.environ["GLYPHSCOPE_ALLOW_MACROS"]

    # Check macro NOT used
    if trace.get("macro_used") is not None:
        print(f"❌ FAIL: Macro used despite GLYPHSCOPE_ALLOW_MACROS=0")
        return False

    if trace.get("macros_allowed") != False:
        print(f"❌ FAIL: macros_allowed should be False")
        return False

    # Check body_hint applied
    if not trace.get("body_scaffold_applied"):
        print(f"❌ FAIL: Body hint not applied")
        return False

    if trace.get("body_scaffold_tokens_used") != "4/4":
        print(f"❌ FAIL: Body hint incomplete: {trace.get('body_scaffold_tokens_used')}")
        return False

    # Check compiles
    if not trace.get("compiled"):
        print(f"❌ FAIL: Code does not compile")
        print(f"Error: {trace.get('compile_error')}")
        return False

    # Check AST OK
    if not trace.get("ast_ok"):
        print(f"❌ FAIL: AST not safe: {trace.get('ast_violation')}")
        return False

    # Check function works (baseline: returns 0)
    namespace = {}
    try:
        exec(text, namespace)
        levenshtein = namespace.get("levenshtein")
        if not callable(levenshtein):
            print(f"❌ FAIL: No callable 'levenshtein' function")
            return False

        result = levenshtein("test", "test")
        if result != 0:
            print(f"❌ FAIL: levenshtein('test', 'test') = {result}, expected 0 (baseline)")
            return False
    except Exception as e:
        print(f"❌ FAIL: Runtime error: {e}")
        return False

    # Check determinism in no-macro mode
    os.environ["GLYPHSCOPE_ALLOW_MACROS"] = "0"
    decoder2 = CDCDecoder(seed_path=None)
    text2, trace2 = decoder2.generate_from_anchor("", max_new_tokens=64, temperature=0.3, seed=123)
    del os.environ["GLYPHSCOPE_ALLOW_MACROS"]

    if text != text2:
        print(f"❌ FAIL: No-macro mode not deterministic")
        print(f"Text1 len: {len(text)}, Text2 len: {len(text2)}")
        print(f"Tokens1: {trace['generated_tokens']}, Tokens2: {trace2['generated_tokens']}")
        return False

    print(f"✅ PASS: No-macro mode works correctly")
    print(f"   Macros allowed: False")
    print(f"   Body hint: {trace['body_scaffold_tokens_used']}")
    print(f"   Tokens: {trace['generated_tokens']} (scaffold + body_hint)")
    print(f"   Compiles: {trace['compiled']}")
    print(f"   AST OK: {trace['ast_ok']}")
    print(f"   Deterministic: {text == text2}")
    return True


def main():
    print("╔════════════════════════════════════════════════════════════╗")
    print("║  GLYPHSCOPE VAULT-ANCHORED CDC - ACCEPTANCE KATS          ║")
    print("╚════════════════════════════════════════════════════════════╝\n")

    tests = [
        test_vocab_kat,
        test_header_kat,
        test_state_kat,
        test_macro_kat,
        test_compile_kat,
        test_determinism_kat,
        test_no_macro_kat,
    ]

    results = []
    for test in tests:
        try:
            results.append(test())
        except Exception as e:
            print(f"❌ EXCEPTION: {e}")
            import traceback
            traceback.print_exc()
            results.append(False)

    print("\n" + "="*60)
    passed = sum(results)
    total = len(results)

    if passed == total:
        print(f"✅ ALL {total} KATS PASS")
        return 0
    else:
        print(f"❌ {total - passed}/{total} KATS FAILED")
        return 1


if __name__ == "__main__":
    sys.exit(main())
