#!/usr/bin/env python3
"""
Execution-based evaluation for FSI_ECHO.
Fixes tokenizer decode (adds back spaces stripped during encode), then exec tests.
"""
import sys, os, time, re, json, signal
sys.path.insert(0, '/tmp/fsi_echo')
from fsi_echo import FSIEchoModel, CodeTokenizer, export_gguf
import torch

TIMEOUT = 5

# ============================================================
# FIXED DECODE: adds Python-required whitespace between tokens
# ============================================================
# The tokenizer strips ALL spaces during encode.
# During decode we need to insert spaces where Python needs them.
NO_SPACE_BEFORE = {*'.,:;)]}'}
NO_SPACE_AFTER  = {*'([{'}
ALWAYS_SPACE_BEFORE = {'def', 'return', 'if', 'elif', 'else', 'for', 'while',
                        'try', 'except', 'finally', 'with', 'in', 'not', 'and', 'or',
                        'is', 'from', 'import', 'as', 'pass', 'break', 'continue',
                        'raise', 'yield', 'class', 'lambda'}
ALWAYS_SPACE_AFTER = {*',:='}

def decode_with_spaces(tok, ids):
    """Decode token IDs with proper Python whitespace."""
    parts = [tok.inverse.get(i, '<UNK>') for i in ids]
    result = []
    for i, p in enumerate(parts):
        if p == '\n':
            result.append('\n')
            continue
        if not result:
            result.append(p)
            continue
        prev = result[-1]

        # Never space before newline
        if prev == '\n':
            result.append(p)
            continue

        # Never space between certain chars
        need_space = True

        # No space: prev is open paren/bracket and current is close or comma
        if prev in '([{' and p in ')]},:.':
            need_space = False
        # No space: prev and current are both operators
        elif prev in NO_SPACE_AFTER and p in NO_SPACE_BEFORE:
            need_space = False
        # No space: prev is a newline (indentation)
        elif prev == '\n':
            need_space = False
        # No space between word chars and open paren (function call)
        elif prev.isalnum() and p == '(':
            need_space = False
        # Space before keywords (unless continuous word)
        elif p in ALWAYS_SPACE_BEFORE and prev.isalnum():
            need_space = True
        # No space before punctuation
        elif p in NO_SPACE_BEFORE:
            need_space = False
        # No space after open punctuation
        elif prev in NO_SPACE_AFTER:
            need_space = False
        # No space around dot
        elif prev == '.' or p == '.':
            need_space = False
        # No space between word chars and following
        elif prev.isalnum() and p.isalnum():
            need_space = False
        # No space between digits
        elif prev.isdigit() and p.isdigit():
            need_space = False
        # Always space after comma
        elif prev == ',':
            need_space = True

        if need_space:
            result.append(' ')
        result.append(p)

    return ''.join(result)


def smart_decode(model_out, tok, func_name):
    """Extract valid Python function from model output."""
    generated = model_out['generated']
    # The generated string includes prompt + continuation joined with no spaces
    # We need to extract the full function definition

    # Try to find the function
    cleaned = generated
    # Remove special tokens
    cleaned = re.sub(r'<[^>]+>', '', cleaned)

    # Check if the output already has "def" 
    if 'def' in cleaned:
        # Extract from 'def' to end
        idx = cleaned.index('def')
        # Get raw token IDs from the model output? No, we have decoded text.
        # Just try to rebuild from the decoded text.
        
        # The output might look like: "defadd(a,b):\n    returna+b\n"
        # Need to fix: insert space after 'def', after 'return', etc.
        code = cleaned[idx:]
        
        # Simple fixes:
        code = re.sub(r'\b(def|return|if|elif|else|for|while|try|except|with|in|not|and|or|is|from|import|class|raise|yield|pass|break|continue)\b',
                      lambda m: m.group(1) + ' ', code)
        # Fix comma spacing
        code = re.sub(r',(?!\s|$)', ', ', code)
        # Fix colon spacing (but not '::')
        code = re.sub(r':(?!\s|\n|:)', ': ', code)
        # Fix operator spacing around =, ==, !=, <=, >=, <, >, +, -, *, /, %
        # But be careful not to break string literals
        code = re.sub(r'(?<=[a-zA-Z0-9_)])\s*(==|!=|<=|>=|\*\*|//|<<|>>)\s*(?=[a-zA-Z0-9_(])', r' \1 ', code)
        code = re.sub(r'(?<=[a-zA-Z0-9_)\])\s*([+\-*/%<>]|={1,2}|!)\s*(?=[a-zA-Z0-9_(])', r' \1 ', code)
        # Ensure proper negative number handling (don't break -5)
        # Actually just try to compile it
        
        return code

    return None

# ============================================================
# SAFE EXECUTION
# ============================================================
class TimeoutError(Exception): pass
def timeout_handler(signum, frame):
    raise TimeoutError()

def try_generated(code, test_args, expected, func_name):
    """Try to exec code and check output. Returns (passed, detail)."""
    for attempt in range(3):
        # Attempt progressive fixes
        try:
            compile(code, '<eval>', 'exec')
            namespace = {}
            signal.signal(signal.SIGALRM, timeout_handler)
            signal.alarm(TIMEOUT)
            exec(code, namespace)
            signal.alarm(0)
            if func_name not in namespace:
                # Try with leading underscore or different name
                for k in namespace:
                    if func_name in k:
                        fn = namespace[k]
                        break
                else:
                    return False, f"'{func_name}' not in namespace (has: {list(namespace.keys())[:3]})"
            else:
                fn = namespace[func_name]
            result = fn(*test_args)
            if isinstance(expected, float):
                passed = abs(result - expected) < 1e-6
            else:
                passed = result == expected
            return passed, f"got {repr(result)}, expected {repr(expected)}"
        except TimeoutError:
            return False, "TIMEOUT"
        except SyntaxError as e:
            # Try progressively more aggressive fixes
            if attempt == 0:
                # First attempt: Regen code with space fixes
                code = re.sub(r'\b(def|return|if|elif|else|for|while|try|except|with|in|not|and|or|is)\b',
                             lambda m: m.group(1) + ' ', code)
            elif attempt == 1:
                # Second attempt: Add spaces around all operators
                code = re.sub(r'(?<=[a-zA-Z0-9_)\])\s*([+\-*/%<>=!])\s*(?=[a-zA-Z0-9_(])', r' \1 ', code)
                code = re.sub(r',(?!\s)', ', ', code)
            else:
                return False, f"SYNTAX: {e}"
        except Exception as e:
            return False, f"err: {e}"
    return False, "exhausted fixes"

def normalize_args(args):
    if isinstance(args, str):
        return (args,)
    return tuple(args)

# ============================================================
# TEST SUITE
# ============================================================
TEST_SUITE = {
    "add":        [((3, 4), 7), ((-1, 5), 4), ((0, 0), 0), ((100, 200), 300)],
    "sub":        [((10, 3), 7), ((5, 10), -5), ((0, 0), 0), ((-5, -3), -2)],
    "mul":        [((3, 4), 12), ((-2, 5), -10), ((0, 100), 0), ((7, 7), 49)],
    "div":        [((10, 2), 5.0), ((9, 3), 3.0), ((1, 2), 0.5), ((0, 5), 0.0)],
    "square":     [((5,), 25), ((-3,), 9), ((0,), 0), ((10,), 100)],
    "cube":       [((3,), 27), ((-2,), -8), ((0,), 0), ((5,), 125)],
    "double":     [((3,), 6), ((-5,), -10), ((0,), 0), ((100,), 200)],
    "half":       [((10,), 5.0), ((3,), 1.5), ((0,), 0.0), ((-4,), -2.0)],
    "increment":  [((5,), 6), ((-1,), 0), ((0,), 1), ((99,), 100)],
    "decrement":  [((5,), 4), ((0,), -1), ((-5,), -6), ((1,), 0)],
    "modulo":     [((10, 3), 1), ((8, 4), 0), ((7, 5), 2), ((0, 3), 0)],
    "power":      [((2, 3), 8), ((5, 0), 1), ((3, 2), 9), ((10, 1), 10)],
    "negate":     [((5,), -5), ((-3,), 3), ((0,), 0), ((100,), -100)],
    "absolute":   [((5,), 5), ((-5,), 5), ((0,), 0), ((-123,), 123)],
    "is_even":    [((2,), True), ((3,), False), ((0,), True), ((-4,), True)],
    "is_odd":     [((2,), False), ((3,), True), ((0,), False), ((-5,), True)],
    "is_positive":[((5,), True), ((-5,), False), ((0,), False), ((100,), True)],
    "is_negative":[((5,), False), ((-5,), True), ((0,), False), ((-1,), True)],
    "is_zero":    [((0,), True), ((5,), False), ((0,), True), ((100,), False)],
    "is_greater": [((5, 3), True), ((3, 5), False), ((5, 5), False), ((0, -1), True)],
    "is_less":    [((3, 5), True), ((5, 3), False), ((5, 5), False), ((-1, 0), True)],
    "factorial":  [((0,), 1), ((1,), 1), ((5,), 120), ((10,), 3628800)],
    "fibonacci":  [((0,), 0), ((1,), 1), ((10,), 55), ((20,), 6765)],
    "gcd":        [((12, 8), 4), ((7, 3), 1), ((0, 5), 5), ((18, 12), 6)],
    "is_prime":   [((2,), True), ((4,), False), ((17,), True), ((1,), False)],
    "digit_sum":  [((123,), 6), ((0,), 0), ((999,), 27), ((-5,), 5)],
    "digit_count":[((12345,), 5), ((0,), 1), ((100,), 3), ((-50,), 2)],
    "clamp":      [((5, 1, 10), 5), ((0, 1, 10), 1), ((15, 1, 10), 10), ((-5, 0, 100), 0)],
    "sign":       [((5,), 1), ((-5,), -1), ((0,), 0), ((100,), 1)],
    "max_of_two": [((5, 10), 10), ((-5, -10), -5), ((7, 7), 7), ((0, 100), 100)],
    "min_of_two": [((5, 10), 5), ((-5, -10), -10), ((7, 7), 7), ((100, 0), 0)],
    "first":      [(([1,2,3],), 1), ((["a","b"],), "a"), (([True],), True)],
    "last":       [(([1,2,3],), 3), ((["a","b"],), "b"), (([True],), True)],
    "list_length":[(([1,2,3],), 3), (([],), 0), ((["a"],), 1)],
    "list_sum":   [(([1,2,3],), 6), (([],), 0), (([-5,5],), 0)],
    "list_max":   [(([1,5,3],), 5), (([-1,-5,-3],), -1), (([100],), 100)],
    "list_min":   [(([1,5,3],), 1), (([-1,-5,-3],), -5), (([100],), 100)],
    "uppercase":  [("hello", "HELLO"), ("ABC", "ABC"), ("", ""), ("a1b2", "A1B2")],
    "lowercase":  [("HELLO", "hello"), ("abc", "abc"), ("", ""), ("A1B2", "a1b2")],
    "string_length":[("hello", 5), ("", 0), ("abc123", 6)],
    "classify_number": [((5,), 'positive'), ((-3,), 'negative'), ((0,), 'zero')],
    "sum_to":     [((5,), 15), ((0,), 0), ((1,), 1), ((100,), 5050)],
}

# ============================================================
# MAIN
# ============================================================
device = 'cpu'
model = FSIEchoModel().to(device)
tok = CodeTokenizer()

ckpt = '/tmp/fsi_echo/checkpoints/gold_best.pt'
d = torch.load(ckpt, map_location='cpu', weights_only=True)
model.load_state_dict(d['model'])
step = d.get('step', '?')
vl = d.get('val_loss', '?')
if isinstance(vl, float):
    print(f"Loaded {ckpt} (step {step}, val_loss={vl:.4f})")
else:
    print(f"Loaded {ckpt}")
print(f"Model: {model.param_count():,} params\n")

all_passed = 0
all_total = 0

for func_name in sorted(TEST_SUITE.keys()):
    r = model.generate(tok, "def " + func_name, max_tokens=120, temperature=0.1)
    raw = r['generated']

    # Extract code from raw output
    code = smart_decode(r, tok, func_name)
    if code is None:
        code = f"def {func_name}():\n    pass\n"

    case_results = []
    for args, expected in TEST_SUITE[func_name]:
        args = normalize_args(args)
        passed, detail = try_generated(code, args, expected, func_name)
        case_results.append((args, expected, passed, detail))

    passed = sum(1 for _, _, p, _ in case_results if p)
    total = len(case_results)
    all_passed += passed
    all_total += total

    s = "✓" if passed == total else f"✗({passed}/{total})"
    display = repr(code[:70])
    print(f"  {s} {func_name:<16s} | {display}")

    for args, expected, p, detail in case_results:
        if not p:
            print(f"       args={args} → {detail}")

pct = 100 * all_passed / all_total if all_total else 0
print(f"\n{'='*60}")
print(f"EXECUTION ACCURACY: {all_passed}/{all_total} ({pct:.1f}%)")
print(f"{'='*60}")