#!/usr/bin/env python3
"""Py-Agent HumanEval & Code Quality Benchmark Suite [Fast Local Edition]"""

import ast
import contextlib
import io
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.request as urlreq

CFG_DIR = os.path.expanduser("~/.config/py-agent")
MODULES_DIR = os.path.join(CFG_DIR, "modules")

if MODULES_DIR not in sys.path:
    sys.path.insert(0, MODULES_DIR)
if CFG_DIR not in sys.path:
    sys.path.insert(0, CFG_DIR)

try:
    import agent_core as core
    import agent_cloud as cloud
    import agent_ui as ui
    from rich.console import Console
    from rich.table import Table
    from rich.panel import Panel
except ImportError as e:
    sys.stderr.write(f"Import error: {e}\n")
    sys.exit(1)

console = Console()

# ── Curated High-Speed HumanEval Benchmark Suite ──────────────────────────────
BENCHMARKS = [
    {
        "id": "HE-01",
        "category": "Strings & Stacks",
        "name": "Balanced Brackets",
        "prompt": "Write a python function `is_balanced(s: str) -> bool` that returns True if all parentheses '()', brackets '[]', and braces '{}' in string `s` are correctly balanced, and False otherwise. Output ONLY the python code block.",
        "test_code": """
assert is_balanced("()") == True
assert is_balanced("([{}])") == True
assert is_balanced("([)]") == False
assert is_balanced("{[()]}") == True
assert is_balanced("(((") == False
"""
    },
    {
        "id": "HE-02",
        "category": "Math & Primes",
        "name": "Prime Factorization",
        "prompt": "Write a python function `prime_factors(n: int) -> list[int]` that returns a sorted list of all prime factors of integer `n`. Output ONLY the python code block.",
        "test_code": """
assert prime_factors(100) == [2, 2, 5, 5]
assert prime_factors(37) == [37]
assert prime_factors(84) == [2, 2, 3, 7]
assert prime_factors(2) == [2]
"""
    },
    {
        "id": "HE-03",
        "category": "Data Structures",
        "name": "Frequency Counter",
        "prompt": "Write a python function `top_elements(lst: list, k: int) -> list` that returns the `k` most frequent elements in `lst` sorted in descending order of frequency. Output ONLY the python code block.",
        "test_code": """
assert top_elements([1, 1, 1, 2, 2, 3], 2) == [1, 2]
assert top_elements(["a", "b", "a", "c", "a", "b"], 1) == ["a"]
assert top_elements([4, 4, 4, 4], 1) == [4]
"""
    },
    {
        "id": "HE-04",
        "category": "Algorithms",
        "name": "Binary Search Insert",
        "prompt": "Write a python function `search_insert(nums: list[int], target: int) -> int` that returns the index if the target is found in sorted list `nums`. If not, return the index where it would be if it were inserted in order. Output ONLY the python code block.",
        "test_code": """
assert search_insert([1, 3, 5, 6], 5) == 2
assert search_insert([1, 3, 5, 6], 2) == 1
assert search_insert([1, 3, 5, 6], 7) == 4
assert search_insert([1, 3, 5, 6], 0) == 0
"""
    },
    {
        "id": "HE-05",
        "category": "Recursion",
        "name": "Deep Flatten",
        "prompt": "Write a python function `flatten_deep(nested: list) -> list` that recursively flattens an arbitrarily nested list of integers into a single flat list. Output ONLY the python code block.",
        "test_code": """
assert flatten_deep([1, [2, [3, [4]], 5]]) == [1, 2, 3, 4, 5]
assert flatten_deep([[[1]], 2, [3, 4]]) == [1, 2, 3, 4]
assert flatten_deep([]) == []
"""
    },
    {
        "id": "HE-06",
        "category": "String Parsing",
        "name": "Roman to Integer",
        "prompt": "Write a python function `roman_to_int(s: str) -> int` that converts a Roman numeral string (e.g. 'MCMXCIV') to an integer. Output ONLY the python code block.",
        "test_code": """
assert roman_to_int("III") == 3
assert roman_to_int("LVIII") == 58
assert roman_to_int("MCMXCIV") == 1994
"""
    },
    {
        "id": "HE-07",
        "category": "Logic & Intervals",
        "name": "Merge Intervals",
        "prompt": "Write a python function `merge_intervals(intervals: list[list[int]]) -> list[list[int]]` that merges all overlapping intervals and returns an array of the non-overlapping intervals sorted by start time. Output ONLY the python code block.",
        "test_code": """
assert merge_intervals([[1, 3], [2, 6], [8, 10], [15, 18]]) == [[1, 6], [8, 10], [15, 18]]
assert merge_intervals([[1, 4], [4, 5]]) == [[1, 5]]
assert merge_intervals([[1, 4], [0, 4]]) == [[0, 4]]
"""
    },
    {
        "id": "HE-08",
        "category": "Math & Logic",
        "name": "LCM Calculator",
        "prompt": "Write a python function `lcm(a: int, b: int) -> int` that calculates the least common multiple of two positive integers `a` and `b`. Output ONLY the python code block.",
        "test_code": """
assert lcm(4, 6) == 12
assert lcm(7, 13) == 91
assert lcm(10, 15) == 30
assert lcm(1, 42) == 42
"""
    }
]


def extract_python_code(raw_response: str) -> str:
    """Extracts raw python code from markdown or plain text."""
    clean = re.sub(r"<think>[\s\S]*?</think>", "", raw_response).strip()
    if m := re.search(r"```(?:python)?\s*([\s\S]*?)\s*```", clean, re.IGNORECASE):
        return m.group(1).strip()
    return clean


def run_single_test(bench: dict, model_name: str) -> tuple[bool, float, int, str]:
    """Runs a single HumanEval task against the active model and verifies with deterministic assertions."""
    messages = [
        {"role": "system", "content": "You are an expert Python programming assistant. Write clean, correct, self-contained Python code. Output ONLY the python code block with no conversational text."},
        {"role": "user", "content": bench["prompt"]}
    ]
    
    t0 = time.time()
    # Suppress raw stdout streaming so the progress table renders cleanly
    sink = io.StringIO()
    with contextlib.redirect_stdout(sink), contextlib.redirect_stderr(sink):
        response = core.stream_response(messages, prefix="AI:", show_stats=False, thinking_budget=0)
    elapsed = time.time() - t0
    
    if not response:
        return False, elapsed, 0, "Empty response"
    
    code = extract_python_code(response)
    toks = core.get_accurate_token_count(response)
    
    test_harness = f"{code}\n\n{bench['test_code']}"
    
    sandbox_dir = tempfile.mkdtemp(prefix="he_eval_")
    test_file = os.path.join(sandbox_dir, "test_eval.py")
    
    try:
        with open(test_file, "w", encoding="utf-8") as f:
            f.write(test_harness)
        
        r = subprocess.run([sys.executable, test_file], capture_output=True, text=True, timeout=5)
        passed = (r.returncode == 0)
        err_msg = r.stderr.strip().splitlines()[-1] if r.stderr else ""
        return passed, elapsed, toks, err_msg
    except subprocess.TimeoutExpired:
        return False, elapsed, toks, "Execution Timed Out (>5s)"
    except Exception as e:
        return False, elapsed, toks, str(e)
    finally:
        shutil.rmtree(sandbox_dir, ignore_errors=True)


def get_current_model_info() -> tuple[str, str]:
    configs = cloud.get_active_configs([])
    if configs:
        url = configs[0][0]
        model = configs[0][2].get("model", "cloud-model")
        return model, url
    
    # Check live local server for alias
    local_name = ui.get_local_model_name()
    return local_name, "http://127.0.0.1:8080"


def main():
    console.print(Panel("[bold cyan]Py-Agent HumanEval & Code Quality Benchmark Suite[/bold cyan]\n[dim]Deterministic assertion verification & tokens/sec speed testing[/dim]", border_style="cyan"))
    
    model_name, endpoint = get_current_model_info()
    
    console.print(f"• [bold green]Active Target Model[/bold green]: [bold]{model_name}[/bold]")
    console.print(f"• [dim]Target Endpoint[/dim]: {endpoint}\n")
    
    results = []
    passed_count = 0
    total_tokens = 0
    total_time = 0.0
    
    for idx, bench in enumerate(BENCHMARKS, start=1):
        console.print(f"[dim][{idx}/{len(BENCHMARKS)}][/dim] Evaluating [cyan]{bench['id']}[/cyan] ({bench['name']})...", end="\r")
        passed, elapsed, toks, err = run_single_test(bench, model_name)
        
        total_tokens += toks
        total_time += elapsed
        if passed:
            passed_count += 1
            status_str = "[bold green]PASS[/bold green]"
            detail = f"{elapsed:.2f}s ({toks}t)"
        else:
            status_str = "[bold red]FAIL[/bold red]"
            detail = f"[dim red]{err[:32]}[/dim red]"
            
        tps = (toks / elapsed) if elapsed > 0 else 0
        results.append((bench["id"], bench["name"], bench["category"], status_str, detail, f"{tps:.1f} t/s"))
        
        console.print(f"  {'✔' if passed else '✖'} {status_str} : [bold]{bench['id']}[/bold] {bench['name']:<25} {detail:<22} @ [cyan]{tps:.1f} t/s[/cyan]")
    
    accuracy = (passed_count / len(BENCHMARKS)) * 100
    avg_tps = (total_tokens / total_time) if total_time > 0 else 0
    
    # Leaderboard Summary Table
    table = Table(title=f"\nBenchmark Results: {model_name}", border_style="bright_black")
    table.add_column("ID", style="cyan", justify="left")
    table.add_column("Problem", style="white")
    table.add_column("Domain", style="dim")
    table.add_column("Result", justify="center")
    table.add_column("Details", style="dim")
    table.add_column("Speed", style="green", justify="right")
    
    for row in results:
        table.add_row(*row)
        
    console.print(table)
    
    score_color = "green" if accuracy >= 80 else "yellow" if accuracy >= 50 else "red"
    console.print(f"\n[bold]Final Scorecard:[/bold] [{score_color}]{passed_count}/{len(BENCHMARKS)} Passed ({accuracy:.1f}%)[/{score_color}] | Avg Speed: [bold cyan]{avg_tps:.1f} t/s[/bold cyan] | Total Time: [dim]{total_time:.2f}s[/dim]\n")


if __name__ == "__main__":
    main()
