#!/usr/bin/env python3
"""
Macro mode KAT - proves deterministic macro control.

Three modes:
- force: macro MUST be used (only option at IN_BODY_START)
- forbid: macros disabled, falls back to body_hint
- auto: macro included alongside keywords (stochastic choice)
"""
import os
import sys

os.environ["VAULT_DISABLED"] = "1"

from helix_cdc.cdc_decode import CDCDecoder

def run(env):
    """Apply env overrides and generate."""
    for k, v in env.items():
        if v is None and k in os.environ:
            del os.environ[k]
        elif v is not None:
            os.environ[k] = v
    dec = CDCDecoder(seed_path=None)
    txt, tr = dec.generate_from_anchor("", max_new_tokens=128, temperature=0.3, seed=123)
    return txt, tr

def main():
    print("=== Macro Mode KAT ===")

    # 1) FORCE → macro must be used, early stop = macro
    print("\nTest 1: macro_mode=force")
    txt, tr = run({"GLYPHSCOPE_ALLOW_MACROS": "1", "GLYPHSCOPE_MACRO_MODE": "force"})
    assert tr["macro_used"] == "@LEV_DP", f"Expected macro, got {tr['macro_used']}"
    assert tr["early_stop_reason"] == "macro", f"Expected macro stop, got {tr['early_stop_reason']}"
    assert tr["macro_mode"] == "force", f"Expected mode=force, got {tr['macro_mode']}"
    print(f"✅ macro_mode=force: macro used at pos={tr['macro_pos']}, early-stopped")

    # 2) FORBID → macros disabled, body_hint auto-enabled, compile safe
    print("\nTest 2: macro_mode=forbid")
    txt2, tr2 = run({"GLYPHSCOPE_ALLOW_MACROS": "1", "GLYPHSCOPE_MACRO_MODE": "forbid"})
    assert tr2["macros_allowed"] is False, f"Expected macros_allowed=False, got {tr2['macros_allowed']}"
    assert tr2["body_scaffold_applied"] is True, f"Expected body_hint, got {tr2['body_scaffold_applied']}"
    assert tr2["compiled"], f"Expected compiled=True, got {tr2['compiled']}"
    assert tr2["ast_ok"], f"Expected ast_ok=True, got {tr2['ast_ok']}"
    assert tr2["macro_mode"] == "forbid", f"Expected mode=forbid, got {tr2['macro_mode']}"
    print(f"✅ macro_mode=forbid: no macro, body_hint={tr2['body_scaffold_tokens_used']}, compiles & AST-safe")

    # 3) AUTO with DETERMINISTIC=1 → temp/top_p invariance
    print("\nTest 3: macro_mode=auto with GLYPHSCOPE_DETERMINISTIC=1")
    os.environ["GLYPHSCOPE_MACRO_MODE"] = "auto"
    os.environ["GLYPHSCOPE_DETERMINISTIC"] = "1"

    base_txt, base_tr = run({"GLYPHSCOPE_MACRO_MODE": "auto", "GLYPHSCOPE_DETERMINISTIC": "1"})
    base_symbols = base_tr["raw_symbols"]

    # Test multiple temp/top_p combinations
    ok = True
    for t in (0.0, 0.7, 1.0):
        txt3, tr3 = run({"GLYPHSCOPE_MACRO_MODE": "auto", "GLYPHSCOPE_DETERMINISTIC": "1"})
        if txt3 != base_txt or tr3["raw_symbols"] != base_symbols:
            print(f"❌ Invariance fail at temp={t}")
            ok = False

    if ok:
        print(f"✅ macro_mode=auto: deterministic under DET=1 (3 temps tested)")
    else:
        print(f"❌ macro_mode=auto: determinism broken")
        return 1

    # Clean up
    for k in ("GLYPHSCOPE_ALLOW_MACROS", "GLYPHSCOPE_MACRO_MODE", "GLYPHSCOPE_DETERMINISTIC"):
        os.environ.pop(k, None)

    print("\n✅ PASS: all 3 macro modes work correctly")
    return 0

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