#!/usr/bin/env python3
"""Py-Agent Full-Stack Agentic & Tool-Calling Benchmark Suite [Decision-Grade Edition]"""

import contextlib
import io
import json
import os
import re
import select
import shutil
import subprocess
import sys
import tempfile
import termios
import threading
import time
import tty

CFG_DIR = os.path.expanduser("~/.config/py-agent")
MODULES_DIR = os.path.join(CFG_DIR, "modules")
RESULTS_LOG = os.path.join(CFG_DIR, "tools", "evals", "eval_history.json")
FAILURES_LOG = os.path.join(CFG_DIR, "tools", "evals", "last_failure_trace.log")

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_adapters as adapters
    import agent_cloud as cloud
    import agent_core as core
    import agent_ipython as ipython
    import agent_skills as skills
    import agent_tools as tools
    import agent_ui as ui
    from rich.box import ROUNDED
    from rich.console import Console
    from rich.panel import Panel
    from rich.table import Table
    from rich.text import Text
except ImportError as e:
    sys.stderr.write(f"Import error: {e}\n")
    sys.exit(1)

console = Console()


class MinimalSpinner:
    def __init__(self):
        self.label = "Working..."
        self.running = False
        self.thread = None
        self.chars = ["-", "\\", "|", "/"]
        self.target_stream = getattr(sys, "__stderr__", sys.stderr)

    def start(self, label: str = "Working..."):
        self.label = label
        self.running = True
        self.start_time = time.time()
        self.thread = threading.Thread(target=self._spin, daemon=True)
        self.thread.start()

    def _spin(self):
        i = 0
        try:
            self.target_stream.write("\033[?25l")
            self.target_stream.flush()
        except OSError:
            pass

        while self.running:
            elapsed = time.time() - self.start_time
            try:
                self.target_stream.write(
                    f"\r\033[1;36m[{self.chars[i % len(self.chars)]}]\033[0m {self.label} \033[2m{elapsed:.1f}s\033[0m"
                )
                self.target_stream.flush()
            except OSError:
                pass
            i += 1
            time.sleep(0.08)

    def stop(self):
        self.running = False
        if self.thread and self.thread.is_alive():
            self.thread.join(timeout=0.2)
        try:
            self.target_stream.write("\r\033[2K\r\033[?25h")
            self.target_stream.flush()
        except OSError:
            pass


def _safe_read(path: str) -> str:
    if os.path.isfile(path):
        try:
            with open(path, "r", encoding="utf-8", errors="ignore") as f:
                return f.read()
        except OSError:
            pass
    return ""


AGENTIC_BENCHMARKS = [
    {
        "id": "AG-01",
        "name": "In-Memory Math & Exit [IPython]",
        "tools_tested": "exec_python, final_answer",
        "requires_ipython": True,
        "par_turns": 2,
        "prompt": "Calculate the sum of all multiples of 3 or 5 below 100 in-memory using exec_python. Call final_answer with the integer result and finish.",
        "verify_fn": lambda ws, out: "2318" in str(out) or "2318" in str(getattr(ipython, "_final_answer_val", ""))
    },
    {
        "id": "AG-02",
        "name": "File Create & Overwrite",
        "tools_tested": "write_file (overwrite)",
        "par_turns": 4,
        "prompt": "1. Use write_file to create 'config.json' with {\"version\": 1}.\n2. Rewrite 'config.json' with overwrite=true to {\"version\": 2, \"status\": \"ready\"}.\n3. Conclude immediately with your task complete marker.",
        "verify_fn": lambda ws, out: (
            os.path.isfile(os.path.join(ws, "config.json"))
            and json.loads(_safe_read(os.path.join(ws, "config.json")) or "{}").get("version") == 2
        )
    },
    {
        "id": "AG-03",
        "name": "Surgical Edit & Test",
        "tools_tested": "read_file, edit_file, run_command",
        "par_turns": 6,
        "setup_fn": lambda ws: open(os.path.join(ws, "calc.py"), "w").write(
            "def multiply(a, b):\n    return a + b\n\nif __name__ == '__main__':\n    assert multiply(3, 4) == 12, 'Test failed'\n    print('PASSED')\n"
        ),
        "prompt": "The file 'calc.py' has a bug where multiply(a, b) returns a + b instead of a * b. Inspect it with read_file, fix the formula using edit_file, verify it with run_command('python calc.py'), and finish with your task complete marker.",
        "verify_fn": lambda ws, out: (
            "return a * b" in _safe_read(os.path.join(ws, "calc.py"))
            and subprocess.run([sys.executable, os.path.join(ws, "calc.py")], cwd=ws, capture_output=True).returncode == 0
        )
    },
    {
        "id": "AG-04",
        "name": "Code Search & Multi-File",
        "tools_tested": "search_code, edit_file",
        "par_turns": 8,
        "setup_fn": lambda ws: (
            open(os.path.join(ws, "mod_a.py"), "w").write("API_KEY = 'OLD_KEY'\n"),
            open(os.path.join(ws, "mod_b.py"), "w").write("TARGET = 'OLD_KEY'\n")
        ),
        "prompt": "Search for 'OLD_KEY' across all files using search_code. Replace 'OLD_KEY' with 'NEW_KEY' in every matched file using edit_file. Finish with your task complete marker.",
        "verify_fn": lambda ws, out: (
            "NEW_KEY" in _safe_read(os.path.join(ws, "mod_a.py"))
            and "NEW_KEY" in _safe_read(os.path.join(ws, "mod_b.py"))
        )
    },
    {
        "id": "AG-05",
        "name": "One-and-Done Discipline",
        "tools_tested": "write_file, exit_discipline",
        "par_turns": 2,
        "prompt": "Create a file named 'done.txt' containing 'PROCESSED'. Once written, STOP immediately and output: 'Task complete: done.txt created'",
        "verify_fn": lambda ws, out: (
            os.path.isfile(os.path.join(ws, "done.txt"))
            and "PROCESSED" in _safe_read(os.path.join(ws, "done.txt"))
        )
    },
    {
        "id": "AG-06",
        "name": "Whitespace-Tolerant Edit",
        "tools_tested": "edit_file (resilient match)",
        "par_turns": 6,
        "setup_fn": lambda ws: open(os.path.join(ws, "format_test.py"), "w").write(
            "def greet(name):\n  # 2-space indentation\n  msg = f'Hello, {name}'\n  return msg\n"
        ),
        "prompt": "In 'format_test.py', replace `msg = f'Hello, {name}'` with `msg = f'Welcome, {name}!'` using edit_file. Preserve file syntax and complete the task.",
        "verify_fn": lambda ws, out: "Welcome, {name}!" in _safe_read(os.path.join(ws, "format_test.py"))
    },
    {
        "id": "AG-07",
        "name": "In-Memory Batch Loop [IPython]",
        "tools_tested": "exec_python, final_answer",
        "requires_ipython": True,
        "par_turns": 2,
        "setup_fn": lambda ws: (
            open(os.path.join(ws, "a.txt"), "w").write("10\n"),
            open(os.path.join(ws, "b.txt"), "w").write("20\n"),
            open(os.path.join(ws, "c.txt"), "w").write("30\n")
        ),
        "prompt": "Using exec_python, write a single Python loop to read 'a.txt', 'b.txt', and 'c.txt', sum their integer contents, and return the total (60) via final_answer(60).",
        "verify_fn": lambda ws, out: "60" in str(out) or "60" in str(getattr(ipython, "_final_answer_val", ""))
    }
]


def list_available_profiles() -> list[tuple[str, str]]:
    skills_dir = os.path.join(CFG_DIR, "skills")
    profiles_dir = os.path.join(skills_dir, "profiles")
    found = []

    if os.path.isdir(profiles_dir):
        for root, _, files in os.walk(profiles_dir):
            for f in sorted(files):
                if f.endswith(".md"):
                    full = os.path.join(root, f)
                    rel = os.path.relpath(full, profiles_dir)
                    prof_key = os.path.splitext(rel)[0]
                    lbl = f"Custom {os.path.basename(prof_key).title()}" if prof_key.startswith("custom/") else prof_key.replace("/pro", " Pro").title()
                    found.append((prof_key, lbl))
    return found


def select_profile_interactive(profiles: list[tuple[str, str]], default_key: str = "custom/base") -> str:
    current_idx = next((i for i, (k, _) in enumerate(profiles) if k == default_key or k in ("custom/lingtiny", "custom/nexn25")), 0)
    num_opts = len(profiles)
    last_lines = 0

    try:
        tty_in = open("/dev/tty", "r")
        tty_fd = tty_in.fileno()
    except Exception:
        tty_in = None
        tty_fd = sys.stdin.fileno() if sys.stdin.isatty() else None

    def _get_k():
        if tty_fd is not None:
            try:
                old = termios.tcgetattr(tty_fd)
                try:
                    tty.setraw(tty_fd)
                    termios.tcflush(tty_fd, termios.TCIFLUSH)
                    b = os.read(tty_fd, 1)
                    if b == b"\x1b" and select.select([tty_fd], [], [], 0.05)[0]:
                        b += os.read(tty_fd, 2)
                    return b.decode("utf-8", errors="ignore")
                finally:
                    termios.tcsetattr(tty_fd, termios.TCSADRAIN, old)
            except Exception:
                pass
        return ""

    try:
        sys.stderr.write("\033[?25l")
        while True:
            if last_lines > 0:
                sys.stderr.write(f"\033[{last_lines}A\r\033[J")

            lines_count = 2
            sys.stderr.write("\n\033[1;36mSelect Profile to Benchmark:\033[0m \033[2m(Up/Down navigate, 1-9 select, Enter run, Esc default)\033[0m\n")

            for idx, (k, lbl) in enumerate(profiles):
                marker = " *" if k in (default_key, "custom/lingtiny", "custom/nexn25") else ""
                prefix = f"\033[1;33m  > {idx+1:2d}. {lbl:<22}\033[0m \033[90m({k})\033[0m{marker}\n" if idx == current_idx else f"\033[37m    {idx+1:2d}. {lbl:<22}\033[0m \033[2m({k})\033[0m{marker}\n"
                sys.stderr.write(prefix)
                lines_count += 1

            sys.stderr.write(f"\n\033[2m:: Enter run [{profiles[current_idx][1]}]  Esc: default\033[0m\n")
            lines_count += 2
            sys.stderr.flush()

            last_lines = lines_count
            key = _get_k()

            if key in ("\x1b[A", "k", "K"):
                current_idx = (current_idx - 1 + num_opts) % num_opts
            elif key in ("\x1b[B", "j", "J"):
                current_idx = (current_idx + 1) % num_opts
            elif key.isdigit() and 1 <= int(key) <= num_opts:
                current_idx = int(key) - 1
                break
            elif key in ("\r", "\n", ""):
                break
            elif key in ("\x03", "\x1b", "q", "Q"):
                current_idx = next((i for i, (k, _) in enumerate(profiles) if k == default_key), 0)
                break

        sys.stderr.write(f"\033[{last_lines}A\r\033[J")
        sys.stderr.flush()
        return profiles[current_idx][0]
    finally:
        sys.stderr.write("\033[?25h")
        sys.stderr.flush()
        if tty_in:
            try:
                tty_in.close()
            except Exception:
                pass


def resolve_dynamic_profile(cli_profile_arg: str | None = None) -> tuple[str, str]:
    skills_dir = os.path.join(CFG_DIR, "skills")
    if cli_profile_arg and not cli_profile_arg.startswith("-"):
        clean_arg = cli_profile_arg.lstrip("-")
        if content := skills.load_skill_content(clean_arg, skills_dir, CFG_DIR):
            return clean_arg, content
        if not clean_arg.startswith("custom/") and (content := skills.load_skill_content(f"custom/{clean_arg}", skills_dir, CFG_DIR)):
            return f"custom/{clean_arg}", content

    if active_env := os.environ.get("AI_ACTIVE_SKILL"):
        first_skill = active_env.split()[0]
        if content := skills.load_skill_content(first_skill, skills_dir, CFG_DIR):
            return first_skill, content

    ws_path = os.environ.get("AI_WORKSPACE_PATH", os.getcwd())
    cfg_file = os.path.join(ws_path, ".agent", "config.json")
    if os.path.isfile(cfg_file):
        try:
            with open(cfg_file, "r", encoding="utf-8") as cf:
                if saved_prof := json.load(cf).get("profile", ""):
                    if content := skills.load_skill_content(saved_prof, skills_dir, CFG_DIR):
                        return saved_prof, content
        except Exception:
            pass

    profiles = list_available_profiles()
    if profiles and not os.environ.get("AI_CONTEXT_RUN"):
        chosen_key = select_profile_interactive(profiles, "custom/base")
        if content := skills.load_skill_content(chosen_key, skills_dir, CFG_DIR):
            return chosen_key, content

    default_prof = "custom/base"
    content = skills.load_skill_content(default_prof, skills_dir, CFG_DIR) or "You are an autonomous AI software engineer."
    return default_prof, content


def get_current_model_info() -> tuple[str, str]:
    if configs := cloud.get_active_configs([]):
        return configs[0][2].get("model", "cloud-model"), configs[0][0]
    return ui.get_local_model_name(), "http://127.0.0.1:8080"


def _log_failure_trace(bench: dict, model_name: str, profile_name: str, messages: list, captured_stdout: str, sandbox_dir: str):
    try:
        os.makedirs(os.path.dirname(FAILURES_LOG), exist_ok=True)
        with open(FAILURES_LOG, "a", encoding="utf-8") as f:
            f.write(f"\n{'='*80}\n")
            f.write(f"FAILED TEST: {bench['id']} - {bench['name']}\n")
            f.write(f"MODEL: {model_name} | PROFILE: {profile_name}\n")
            f.write(f"TIMESTAMP: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"{'='*80}\n\n")
            f.write("--- MESSAGE / TOOL CALL HISTORY ---\n")
            for m in messages:
                role = m.get("role", "unknown").upper()
                content = m.get("content", "")
                f.write(f"[{role}]:\n{content}\n\n")
            f.write("--- CAPTURED AGENT TURNS & STDERR ---\n")
            f.write(captured_stdout + "\n")
    except Exception:
        pass


def run_single_agentic_test(bench: dict, model_name: str, profile_prompt: str) -> tuple[bool, float, int, int, str]:
    orig_cwd = os.getcwd()
    sandbox_dir = tempfile.mkdtemp(prefix="agent_eval_")
    orig_ws = os.environ.get("AI_WORKSPACE_PATH")
    os.environ["AI_WORKSPACE_PATH"] = sandbox_dir
    os.environ["AI_CONFIRM_GATES"] = "0"

    tools.clear_session_tracking()

    if "setup_fn" in bench:
        bench["setup_fn"](sandbox_dir)

    messages = [
        {"role": "system", "content": profile_prompt},
        {"role": "user", "content": bench["prompt"]}
    ]

    configs = cloud.get_active_configs(messages)
    if not configs:
        configs = [("http://localhost:8080/v1/chat/completions", {}, {"messages": messages, "stream": True}, 180)]
    url, headers, body, timeout = configs[0]

    t0 = time.time()
    sink = io.StringIO()
    try:
        with contextlib.redirect_stdout(sink), contextlib.redirect_stderr(sink):
            response = core.agentic_turn(
                messages,
                url,
                headers,
                body,
                timeout,
                spinner=None,
                show_stats=False,
                is_agent=True,
            )
        elapsed = time.time() - t0
        turn_count = len([m for m in messages if m.get("role") in ("assistant", "tool")])
        total_tokens = sum(core.get_accurate_token_count(m.get("content") or "") for m in messages)
        passed = bench["verify_fn"](sandbox_dir, response or "")
        if not passed:
            _log_failure_trace(bench, model_name, profile_prompt, messages, sink.getvalue(), sandbox_dir)
        par_val = bench.get("par_turns", turn_count)
        detail = f"{elapsed:.2f}s ({turn_count}/{par_val}t)" if passed else "State verification failed"
        return passed, elapsed, turn_count, total_tokens, detail
    except Exception as e:
        return False, time.time() - t0, 0, 0, f"Exception: {str(e)[:30]}"
    finally:
        os.chdir(orig_cwd)
        shutil.rmtree(sandbox_dir, ignore_errors=True)
        if orig_ws:
            os.environ["AI_WORKSPACE_PATH"] = orig_ws
        else:
            os.environ.pop("AI_WORKSPACE_PATH", None)


def main():
    cli_arg = sys.argv[1] if len(sys.argv) > 1 and not sys.argv[1].startswith("-") else None
    prof_name, profile_content = resolve_dynamic_profile(cli_arg)

    console.print(Panel(
        "[bold cyan]Py-Agent Full-Stack Agentic & Tool-Calling Benchmark Suite[/bold cyan]\n"
        "[dim]Decision-grade evaluation: Par efficiency, token economy, tool resilience, and execution speed[/dim]",
        border_style="cyan"
    ))

    if os.path.isfile(FAILURES_LOG):
        with contextlib.suppress(OSError):
            os.remove(FAILURES_LOG)

    model_name, endpoint = get_current_model_info()
    is_py_enabled = bool(core.get_state("ipython_mode", False))

    mode_label = "Dual (Python + Native)" if is_py_enabled else "Native 6-Tool"
    console.print(f"Target Model: [bold]{model_name}[/bold]")
    console.print(f"Target Endpoint: {endpoint}")
    console.print(f"Harness Profile: [cyan]{prof_name}[/cyan] ({mode_label})\n")

    results = []
    passed_count = 0
    total_time = 0.0
    total_tokens = 0
    total_actual_turns = 0
    total_par_turns = 0
    evaluated_count = 0
    spinner = MinimalSpinner()

    try:
        for idx, bench in enumerate(AGENTIC_BENCHMARKS, start=1):
            if bench.get("requires_ipython") and not is_py_enabled:
                results.append((bench["id"], bench["name"], bench["tools_tested"], "[yellow]SKIP[/yellow]", "ipython: OFF", "-", "-"))
                console.print(f"  [-] SKIP : [bold]{bench['id']}[/bold] {bench['name']:<33} (Skipped: /py disabled)")
                continue

            evaluated_count += 1
            par = bench.get("par_turns", 1)
            total_par_turns += par

            spinner.start(f"[{idx}/{len(AGENTIC_BENCHMARKS)}] Running {bench['id']} ({bench['name']})...")
            passed, elapsed, turns, toks, detail = run_single_agentic_test(bench, model_name, profile_content)
            spinner.stop()

            total_time += elapsed
            total_tokens += toks
            total_actual_turns += turns

            if passed:
                passed_count += 1
                status_str = "[bold green]PASS[/bold green]"
            else:
                status_str = "[bold red]FAIL[/bold red]"

            tps = (toks / elapsed) if elapsed > 0 else 0
            turn_metrics = f"{turns}t (Par {par})" if passed else "FAILED"
            results.append((bench["id"], bench["name"], bench["tools_tested"], status_str, turn_metrics, f"{toks:,}t", f"{tps:.1f} t/s"))
            console.print(f"  {'[green]OK[/green]' if passed else '[red]XX[/red]'} {status_str} : [bold]{bench['id']}[/bold] {bench['name']:<33} {detail:<20} @ [cyan]{tps:.1f} t/s[/cyan]")
    finally:
        spinner.stop()

    accuracy = (passed_count / evaluated_count * 100) if evaluated_count > 0 else 0
    par_eff = (total_par_turns / total_actual_turns * 100) if total_actual_turns > 0 else 0
    avg_tps = (total_tokens / total_time) if total_time > 0 else 0
    toks_per_turn = (total_tokens / total_actual_turns) if total_actual_turns > 0 else 0

    speed_score = min(100.0, (avg_tps / 50.0) * 100.0)
    agent_index = (accuracy * 0.50) + (min(100.0, par_eff) * 0.35) + (speed_score * 0.15)
    grade = "A+" if agent_index >= 95 else "A" if agent_index >= 88 else "B" if agent_index >= 75 else "C" if agent_index >= 60 else "F"

    table = Table(title=f"\nAgentic Stack Benchmark: {model_name} [{prof_name}]", border_style="bright_black")
    table.add_column("ID", style="cyan", justify="left")
    table.add_column("Agent Challenge", style="white")
    table.add_column("Tools Exercised", style="dim")
    table.add_column("Result", justify="center")
    table.add_column("Turns / Par", justify="center")
    table.add_column("Context", justify="right")
    table.add_column("Speed", style="green", justify="right")

    for row in results:
        table.add_row(*row)

    console.print(table)

    summary_text = Text()
    summary_text.append(f"Pass Rate:        {passed_count}/{evaluated_count} ({accuracy:.1f}%)\n", style="bold green" if accuracy >= 80 else "bold red")
    summary_text.append(f"Par Efficiency:   {par_eff:.1f}% ({total_actual_turns} actual turns / {total_par_turns} par target)\n", style="cyan")
    summary_text.append(f"Token Economy:    {total_tokens:,} total tokens (~{toks_per_turn:.0f} tokens/turn)\n", style="dim")
    summary_text.append(f"Effective Speed:  {avg_tps:.1f} tokens/second (Total elapsed: {total_time:.2f}s)\n", style="dim")
    summary_text.append(f"\nAgent Performance Index: {agent_index:.1f}/100 [Grade: {grade}]", style="bold yellow")

    console.print(Panel(summary_text, title="Conclusive Benchmark Evaluation", border_style="green" if accuracy == 100 else "yellow", box=ROUNDED))

    try:
        os.makedirs(os.path.dirname(RESULTS_LOG), exist_ok=True)
        record = {
            "timestamp": int(time.time()),
            "model": model_name,
            "profile": prof_name,
            "accuracy": accuracy,
            "par_efficiency": par_eff,
            "total_tokens": total_tokens,
            "tokens_per_turn": toks_per_turn,
            "avg_tps": avg_tps,
            "total_time": total_time,
            "agent_index": agent_index,
            "grade": grade,
        }
        history = []
        if os.path.isfile(RESULTS_LOG):
            with open(RESULTS_LOG, "r") as rf:
                history = json.load(rf)
        history.append(record)
        with open(RESULTS_LOG, "w") as wf:
            json.dump(history[-30:], wf, indent=2)
    except Exception:
        pass


if __name__ == "__main__":
    main()
