#!/usr/bin/env python3
"""
Early-stop integrity KAT.

Proves: early_stop_reason accurately reflects why generation stopped.
Three cases: macro, body_hint, max_tokens.
"""
import os
import sys

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

from helix_cdc.cdc_decode import CDCDecoder

def main():
    print("=== Early-Stop Integrity KAT ===")

    # Test 1: Stop on macro
    print("\nTest 1: early_stop_reason = 'macro'")
    os.environ["GLYPHSCOPE_MACRO_MODE"] = "force"
    dec1 = CDCDecoder(seed_path=None)
    _, tr1 = dec1.generate_from_anchor("", max_new_tokens=128, temperature=0.2, seed=123)

    assert tr1["early_stop_reason"] == "macro", f"Expected 'macro', got {tr1['early_stop_reason']}"
    assert tr1["macro_used"] == "@LEV_DP", f"Expected macro used"
    assert tr1["macro_pos"] is not None, f"Expected macro_pos"
    print(f"✅ Stopped on macro at pos={tr1['macro_pos']}, tiles={tr1['tiles_before_macro']}")

    # Test 2: Stop on body_hint
    print("\nTest 2: early_stop_reason = 'body_hint'")
    os.environ["GLYPHSCOPE_MACRO_MODE"] = "forbid"
    dec2 = CDCDecoder(seed_path=None)
    _, tr2 = dec2.generate_from_anchor("", max_new_tokens=128, temperature=0.2, seed=123)

    assert tr2["early_stop_reason"] == "body_hint", f"Expected 'body_hint', got {tr2['early_stop_reason']}"
    assert tr2["body_scaffold_applied"] is True, f"Expected body_hint applied"
    assert tr2["macro_used"] is None, f"Expected no macro"
    print(f"✅ Stopped on body_hint after {tr2['body_scaffold_tokens_used']} scaffold tokens")

    # Test 3: Stop on max_tokens (no macro, no body_hint)
    print("\nTest 3: early_stop_reason = 'max_tokens'")
    os.environ["GLYPHSCOPE_MACRO_MODE"] = "auto"
    os.environ["GLYPHSCOPE_MACRO_QUOTA"] = "0"  # No macros allowed
    os.environ["GLYPHSCOPE_BODY_HINT"] = "0"  # Disable body hint
    dec3 = CDCDecoder(seed_path=None)
    _, tr3 = dec3.generate_from_anchor("", max_new_tokens=32, temperature=0.2, seed=123)

    assert tr3["early_stop_reason"] == "max_tokens", f"Expected 'max_tokens', got {tr3['early_stop_reason']}"
    assert tr3["generated_tokens"] == 32, f"Expected exactly 32 tokens, got {tr3['generated_tokens']}"
    assert tr3["macro_used"] is None, f"Expected no macro"
    assert tr3["body_scaffold_applied"] is False, f"Expected no body_hint"
    print(f"✅ Stopped on max_tokens after {tr3['generated_tokens']} tokens")

    # Clean up
    os.environ.pop("GLYPHSCOPE_MACRO_MODE", None)
    os.environ.pop("GLYPHSCOPE_MACRO_QUOTA", None)
    os.environ.pop("GLYPHSCOPE_BODY_HINT", None)

    print("\n✅ PASS: all 3 early-stop reasons work correctly")
    return 0

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