"""
test_battle.py — batch, objective comparison of the fine-tuned LoRA vs base.

Run this on the same machine/Space as app.py (same directory). It imports
the already-fixed model loading + prompt-building logic from app.py, so
there is exactly one source of truth for how prompts are built.

What this adds over the Gradio UI:
  - Runs a fixed battery of test cases across all 4 trained task types.
  - For tasks with an objectively checkable answer, it actually EXECS the
    generated code and asserts against expected output. Syntax-valid code
    that returns the wrong answer will be caught here — the Gradio app's
    heuristic score can't tell you that.
  - Averages results across the batch so one lucky/unlucky prompt doesn't
    decide the verdict.

Usage:
    python test_battle.py
"""

import io
import contextlib
import traceback

# Reuses model loading + the corrected prompt builders from app.py.
# Importing app.py loads the model once (takes a minute); it will NOT
# launch the Gradio UI because that's gated behind `if __name__ == "__main__"`.
import app


# ----------------------------------------------------------------------------
# TEST CASES
# Each has: task type, instruction, and an optional `verify(code)` callable
# that execs the generated code and returns (passed: bool, detail: str).
# No verifier => structural/heuristic scoring only (REFACTOR, CODE_REVIEW,
# and any open-ended DEBUG explanation don't have one "correct" rewrite).
# ----------------------------------------------------------------------------

def _exec_and_get(code: str, names: list[str]):
    """Exec code in an isolated namespace, return the requested names."""
    ns = {}
    exec(code, ns)
    return [ns[n] for n in names]


def verify_palindrome(code: str):
    try:
        (fn,) = _exec_and_get(code, ["is_palindrome"])
        cases = [
            ("A man, a plan, a canal: Panama", True),
            ("hello", False),
            ("", True),
            ("No lemon, no melon", True),
        ]
        for s, expected in cases:
            if fn(s) != expected:
                return False, f"is_palindrome({s!r}) = {fn(s)!r}, expected {expected!r}"
        return True, "all cases passed"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


def verify_fibonacci(code: str):
    try:
        (fn,) = _exec_and_get(code, ["fibonacci"])
        cases = [(0, 0), (1, 1), (5, 5), (10, 55)]
        for n, expected in cases:
            got = fn(n)
            if got != expected:
                return False, f"fibonacci({n}) = {got!r}, expected {expected!r}"
        return True, "all cases passed"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


def verify_merge_sorted(code: str):
    try:
        (fn,) = _exec_and_get(code, ["merge_sorted"])
        got = fn([1, 3, 5], [2, 4, 6])
        if got != [1, 2, 3, 4, 5, 6]:
            return False, f"merge_sorted([1,3,5],[2,4,6]) = {got!r}"
        got2 = fn([], [1, 2])
        if got2 != [1, 2]:
            return False, f"merge_sorted([],[1,2]) = {got2!r}"
        return True, "all cases passed"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


def verify_debug_add(code: str):
    try:
        (fn,) = _exec_and_get(code, ["add"])
        if fn(2, 3) != 5:
            return False, f"add(2, 3) = {fn(2, 3)!r}, expected 5"
        return True, "bug fixed correctly"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


def verify_retry_decorator(code: str):
    try:
        ns = {}
        exec(code, ns)
        decorator_name = next(
            (n for n, v in ns.items() if callable(v) and n.lower().find("retry") != -1),
            None,
        )
        if decorator_name is None:
            return False, "no retry-named callable found in generated code"
        retry = ns[decorator_name]

        attempts = {"n": 0}

        @retry
        def flaky():
            attempts["n"] += 1
            if attempts["n"] < 3:
                raise ValueError("not yet")
            return "ok"

        result = flaky()
        if result != "ok" or attempts["n"] < 3:
            return False, f"expected 3 attempts ending in 'ok', got n={attempts['n']}, result={result!r}"
        return True, f"succeeded after {attempts['n']} attempts"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


TEST_CASES = [
    {
        "task": "GENERATE",
        "instruction": (
            "Write a function called `is_palindrome(s: str) -> bool` that checks "
            "if a string is a valid palindrome, ignoring punctuation, spaces, and case."
        ),
        "verify": verify_palindrome,
    },
    {
        "task": "GENERATE",
        "instruction": (
            "Write a function called `fibonacci(n: int) -> int` that returns the "
            "nth Fibonacci number (0-indexed, fibonacci(0)=0, fibonacci(1)=1) "
            "using iteration, not recursion."
        ),
        "verify": verify_fibonacci,
    },
    {
        "task": "GENERATE",
        "instruction": (
            "Write a function called `merge_sorted(a: list, b: list) -> list` that "
            "merges two already-sorted lists into one sorted list in O(n) time, "
            "without using the built-in sorted() function."
        ),
        "verify": verify_merge_sorted,
    },
    {
        "task": "GENERATE",
        "instruction": (
            "Write a decorator called `retry` that retries a decorated function up "
            "to 3 times if it raises an exception, before letting the final exception "
            "propagate. No external libraries."
        ),
        "verify": verify_retry_decorator,
    },
    {
        "task": "DEBUG",
        "instruction": (
            "def add(a, b):\n    return a + b\n\nprint(add(2))\n\n"
            "# This raises a TypeError. Find and fix the bug. Keep the function name `add`."
        ),
        "verify": verify_debug_add,
    },
    {
        "task": "REFACTOR",
        "instruction": (
            "def f(x):\n"
            "    y=[]\n"
            "    for i in range(len(x)):\n"
            "        if x[i]%2==0:\n"
            "            y.append(x[i])\n"
            "    return y"
        ),
        "verify": None,  # structural quality only — many valid rewrites
    },
    {
        "task": "CODE_REVIEW",
        "instruction": (
            "def get_user(users, id):\n"
            "    for u in users:\n"
            "        if u['id'] == id:\n"
            "            return u\n"
            "\n"
            "def process(users, id):\n"
            "    user = get_user(users, id)\n"
            "    return user['name'].upper()"
        ),
        "verify": None,  # review quality is subjective, no single fixed rewrite
    },
]


# ----------------------------------------------------------------------------
# RUNNER
# ----------------------------------------------------------------------------
def generate_one(instruction: str, task: str, which: str, max_new_tokens: int = 500):
    """which = 'ft' or 'base'"""
    if which == "ft":
        inputs = app.build_inputs_ft(instruction, task)
        text, elapsed, n_tokens = app._run_generate(inputs, max_new_tokens, 0.7, False)
    else:
        inputs = app.build_inputs_base(instruction)
        with app.model.disable_adapter():
            text, elapsed, n_tokens = app._run_generate(inputs, max_new_tokens, 0.7, False)
    return text, elapsed, n_tokens


def run_case(case: dict) -> dict:
    task, instruction, verify = case["task"], case["instruction"], case["verify"]
    result = {"task": task, "instruction": instruction[:60]}

    for which in ("ft", "base"):
        text, elapsed, n_tokens = generate_one(instruction, task, which)
        code = app.extract_code(text)
        metrics = app.analyze_response(text, elapsed, n_tokens)

        if verify is not None:
            with contextlib.redirect_stdout(io.StringIO()):
                try:
                    passed, detail = verify(code)
                except Exception as e:
                    passed, detail = False, f"harness error: {e}"
        else:
            passed, detail = None, "no verifier (structural only)"

        result[which] = {
            "raw": text,
            "code": code,
            "syntax_valid": metrics["syntax_valid"],
            "quality_score": metrics["quality_score"],
            "tokens_per_sec": metrics["tokens_per_sec"],
            "passed": passed,
            "detail": detail,
        }
    return result


def main():
    app._warmup()  # absorb CUDA cold-start once, outside all timed/scored runs

    results = [run_case(c) for c in TEST_CASES]

    print("\n" + "=" * 80)
    print("RESULTS")
    print("=" * 80)

    ft_scores, base_scores = [], []
    ft_pass, base_pass, verifiable = 0, 0, 0

    for r in results:
        print(f"\n[{r['task']}] {r['instruction']}...")
        for which, label in (("ft", "FINE-TUNED"), ("base", "BASE")):
            m = r[which]
            pass_str = (
                "✅ PASS" if m["passed"] is True
                else "❌ FAIL" if m["passed"] is False
                else "—"
            )
            print(
                f"  {label:10s} | syntax={'ok' if m['syntax_valid'] else 'BROKEN':6s} "
                f"| score={m['quality_score']:5.1f}/100 | {pass_str:8s} "
                f"| {m['tokens_per_sec']:.1f} tok/s"
            )
            if m["passed"] is False or (m["passed"] is True and m["detail"]):
                print(f"             detail: {m['detail']}")

            if which == "ft":
                ft_scores.append(m["quality_score"])
                if m["passed"] is True:
                    ft_pass += 1
            else:
                base_scores.append(m["quality_score"])
                if m["passed"] is True:
                    base_pass += 1
        if r["ft"]["passed"] is not None:
            verifiable += 1

    print("\n" + "=" * 80)
    print("SUMMARY")
    print("=" * 80)
    print(f"Avg quality score   — fine-tuned: {sum(ft_scores)/len(ft_scores):.1f}  |  base: {sum(base_scores)/len(base_scores):.1f}")
    if verifiable:
        print(f"Functional pass rate ({verifiable} verifiable tasks) — fine-tuned: {ft_pass}/{verifiable}  |  base: {base_pass}/{verifiable}")
    print("=" * 80)


if __name__ == "__main__":
    main()