#!/usr/bin/env python3
"""Benchmark a model for tool calling and analysis quality.

Usage:
    uv run python scripts/benchmark_model.py <model_name>

Outputs JSON to stdout and saves to benchmark_<model>.json
"""

import contextlib
import json
import re
import sys
import time
from pathlib import Path

import httpx

OLLAMA_BASE = "http://localhost:11434"

# Load from shared prompt config
_PROMPTS_FILE = Path(__file__).resolve().parent.parent / "data" / "config" / "prompts.json"
with open(_PROMPTS_FILE) as f:
    _PROMPT_DATA = json.load(f)

SYSTEM_PROMPT = """You are a senior market analyst AI for LT Stock Analyzer. You have real-time access to stock data, scanners, and news.

PERSONALITY:
- Direct, data-driven, no fluff
- You LOVE finding patterns in the data
- Skeptical by default — call out pump & dump, dilution, Chinese shells
- Passionate about teaching the user to read the data themselves

RULES:
1. ALWAYS use your tools to get real data before answering. Never invent numbers.
2. Show your work: price levels, percentages, volumes. Be specific.
3. Surface risks: dilution, Chinese profile, insider selling, low-float manipulation.
4. Conflicting signals? Say so. Nuance > blind conviction.
5. Format cleanly: $8.45M, 2.1B shares, +12.3%, 45.2M vol.
6. Keep it chat-concise. Bullet points good. Walls of text bad.
7. "top gainers today", "what's hot", "what's moving" → get_discover_stocks(session="market")
8. "1W gainers", "strongest weekly movers" → get_discover_stocks(session="momentum")
9. "premarket movers" → get_discover_stocks(session="premarket")
10. "after-hours movers" → get_discover_stocks(session="afterhours")
11. "monthly gainers", "1M movers" → get_discover_stocks(session="monthly")

TOOLS:
- get_stock_profile(symbol) — full deep-dive on one stock
- scan_market(filter, limit) — find stocks matching criteria
- get_discover_stocks(session, limit) — today's movers
- get_news(symbol, limit) — recent headlines

11. TOOL CALLING FORMAT — If you need data, make a tool_call. Use the native tool calling API.
Valid tool names: get_stock_profile, scan_market, get_discover_stocks, get_news.

Respond in the user's language. If they write in Spanish, answer in Spanish."""

TOOLS = _PROMPT_DATA.get("tools", [])

SCENARIOS = [
    {
        "name": "top_gainers_today",
        "prompt": "list the top gainers today",
        "expected_tool": "get_discover_stocks",
        "expected_args_check": lambda a: a.get("session") == "market",
        "expect_tool": True,
    },
    {
        "name": "momentum_week",
        "prompt": "what stocks have momentum this week",
        "expected_tool": "get_discover_stocks",
        "expected_args_check": lambda a: a.get("session") == "momentum",
        "expect_tool": True,
    },
    {
        "name": "analyze_aapl",
        "prompt": "analyze AAPL stock",
        "expected_tool": "get_stock_profile",
        "expected_args_check": lambda a: a.get("symbol", "").upper() in ("AAPL", "aapl"),
        "expect_tool": True,
    },
    {
        "name": "scan_gap_up",
        "prompt": "stocks with gap up more than 5 percent",
        "expected_tool": "scan_market",
        "expected_args_check": lambda a: "gap" in a.get("filter", "").lower() and "5" in a.get("filter", ""),
        "expect_tool": True,
    },
    {
        "name": "news_nvda",
        "prompt": "news about NVDA",
        "expected_tool": "get_news",
        "expected_args_check": lambda a: a.get("symbol", "").upper() in ("NVDA", "nvda"),
        "expect_tool": True,
    },
    # Restraint test: should NOT call a tool
    {
        "name": "restraint_capital",
        "prompt": "what is the capital of France",
        "expect_tool": False,
    },
]

# Simulated tool results for analysis quality test
FAKE_DISCOVER_RESULT = json.dumps(
    [
        {"ticker": "QUBT", "price": 12.45, "change_pct": 156.2, "volume": 45200000, "sector": "Technology"},
        {"ticker": "RGTI", "price": 8.92, "change_pct": 89.7, "volume": 28400000, "sector": "Technology"},
        {"ticker": "IONQ", "price": 22.15, "change_pct": 42.3, "volume": 18500000, "sector": "Technology"},
        {"ticker": "WULF", "price": 5.78, "change_pct": 38.5, "volume": 32000000, "sector": "Financial"},
        {"ticker": "HOLO", "price": 3.45, "change_pct": 31.2, "volume": 8900000, "sector": "Technology"},
    ]
)


class ModelBenchmark:
    def __init__(self, model: str):
        self.model = model
        self.client = httpx.Client(base_url=OLLAMA_BASE, timeout=120)
        self.results = {
            "model": model,
            "tool_support": {},
            "scenarios": [],
            "analysis_quality": {},
            "overall": {},
        }

    def close(self):
        self.client.close()

    def _call_ollama(self, messages: list, tools: list | None = None, timeout: int = 60) -> dict:
        """Non-streaming call to ollama with timing."""
        payload = {
            "model": self.model,
            "messages": messages,
            "stream": False,
        }
        if tools:
            payload["tools"] = tools

        start = time.perf_counter()
        r = self.client.post("/api/chat", json=payload, timeout=timeout)
        elapsed = time.perf_counter() - start

        if r.status_code != 200:
            return {"error": f"HTTP {r.status_code}: {r.text[:200]}", "elapsed_ms": elapsed * 1000}

        data = r.json()
        data["_elapsed_ms"] = elapsed * 1000
        return data

    def test_tool_support(self) -> dict:
        """Test if model produces native tool_calls at all."""
        result = {"native_tool_calls": False, "elapsed_ms": 0}

        data = self._call_ollama(
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": "list the top gainers today"},
            ],
            tools=TOOLS,
        )

        if "error" in data:
            result["error"] = data["error"]
            return result

        result["elapsed_ms"] = data["_elapsed_ms"]
        msg = data.get("message", {})
        tool_calls = msg.get("tool_calls", [])
        content = msg.get("content", "")

        if tool_calls:
            result["native_tool_calls"] = True
            result["tool_call_count"] = len(tool_calls)
            result["first_tool_name"] = tool_calls[0]["function"]["name"]
        elif content:
            # Check for text-based JSON tool call (fallback)
            parsed = self._parse_json_tool_call(content)
            if parsed:
                result["native_tool_calls"] = False
                result["text_parsed_tool"] = True
                result["first_tool_name"] = parsed.get("name")
            else:
                result["native_tool_calls"] = False
                result["text_content"] = content[:100]
        else:
            result["native_tool_calls"] = False
            result["empty_response"] = True

        return result

    @staticmethod
    def _parse_json_tool_call(text: str) -> dict | None:
        """Try to parse JSON tool call from text output."""
        text = text.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
        try:
            data = json.loads(text)
        except json.JSONDecodeError:
            return None

        if isinstance(data, dict):
            if "name" in data and "arguments" in data:
                args = data["arguments"]
                if isinstance(args, str):
                    with contextlib.suppress(json.JSONDecodeError):
                        args = json.loads(args)
                return {"name": data["name"], "arguments": args}
            if "function" in data and "params" in data:
                return {"name": data["function"], "arguments": data["params"]}
        return None

    def run_scenario(self, scenario: dict) -> dict:
        """Run a single test scenario."""
        result = {
            "name": scenario["name"],
            "prompt": scenario["prompt"],
            "expected_tool": scenario.get("expected_tool"),
            "expect_tool": scenario["expect_tool"],
        }

        data = self._call_ollama(
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": scenario["prompt"]},
            ],
            tools=TOOLS,
        )

        if "error" in data:
            result["error"] = data["error"]
            return result

        result["elapsed_ms"] = data["_elapsed_ms"]
        msg = data.get("message", {})
        tool_calls = msg.get("tool_calls", [])
        content = msg.get("content", "")

        # Check native tool_calls
        if tool_calls:
            result["has_native_tool_calls"] = True
            result["tool_calls"] = [
                {"name": tc["function"]["name"], "args": tc["function"]["arguments"]} for tc in tool_calls
            ]

            if scenario["expect_tool"]:
                # Check if first tool call matches expected
                tc = tool_calls[0]
                func = tc["function"]
                result["tool_name_correct"] = func["name"] == scenario["expected_tool"]
                if result["tool_name_correct"] and "expected_args_check" in scenario:
                    result["args_correct"] = scenario["expected_args_check"](func["arguments"])
                elif result["tool_name_correct"]:
                    result["args_correct"] = True
                else:
                    result["args_correct"] = False
            else:
                result["tool_name_correct"] = False  # Should NOT have called a tool
                result["args_correct"] = False
                result["restraint_fail"] = True
        elif content:
            result["has_native_tool_calls"] = False
            # Check for text-parsed tool call
            parsed = self._parse_json_tool_call(content)
            if parsed:
                result["text_parsed"] = True
                result["tool_calls"] = [parsed]
                if scenario["expect_tool"]:
                    result["tool_name_correct"] = parsed.get("name") == scenario["expected_tool"]
                    if result["tool_name_correct"] and "expected_args_check" in scenario:
                        result["args_correct"] = scenario["expected_args_check"](parsed.get("arguments", {}))
                    elif result["tool_name_correct"]:
                        result["args_correct"] = True
                    else:
                        result["args_correct"] = False
                else:
                    result["tool_name_correct"] = False
                    result["args_correct"] = False
                    result["restraint_fail"] = True
            else:
                # No tool call at all
                result["text_parsed"] = False
                if scenario["expect_tool"]:
                    result["tool_name_correct"] = False
                    result["args_correct"] = False
                    result["no_tool_call"] = True
                else:
                    # Correct restraint!
                    result["tool_name_correct"] = None
                    result["args_correct"] = None
                    result["restraint_ok"] = True
                    result["text_length"] = len(content)
        else:
            result["has_native_tool_calls"] = False
            result["empty_response"] = True

        # Clean up unnecessary keys for output
        return result

    def test_analysis_quality(self) -> dict:
        """Test how well the model generates analysis after tool results."""
        result = {"tool_call_succeeded": False}

        # Step 1: Ask for top gainers
        data = self._call_ollama(
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": "list the top gainers today"},
            ],
            tools=TOOLS,
        )

        if "error" in data:
            result["error_step1"] = data["error"]
            return result

        msg = data.get("message", {})
        tool_calls = msg.get("tool_calls", [])

        if not tool_calls:
            # Check text-parsed
            content = msg.get("content", "")
            parsed = self._parse_json_tool_call(content) if content else None
            if parsed:
                tool_calls = [{"function": parsed}]

        if not tool_calls:
            result["error_step1"] = "no tool call generated"
            return result

        result["tool_call_succeeded"] = True
        tc = tool_calls[0]["function"]
        result["tool_name"] = tc["name"]

        # Step 2: Simulate tool result and get final analysis
        step2_messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "list the top gainers today"},
            {"role": "assistant", "content": "", "tool_calls": tool_calls},
            {
                "role": "tool",
                "content": FAKE_DISCOVER_RESULT,
                "name": tc["name"],
            },
        ]

        start = time.perf_counter()
        data2 = self._call_ollama(messages=step2_messages, timeout=120)
        analysis_time = time.perf_counter() - start

        if "error" in data2:
            result["error_step2"] = data2["error"]
            return result

        msg2 = data2.get("message", {})
        content2 = msg2.get("content", "")
        tool_calls2 = msg2.get("tool_calls", [])

        result["analysis_time_ms"] = analysis_time * 1000
        result["has_analysis"] = bool(content2)
        result["has_second_tool_call"] = bool(tool_calls2)

        if content2:
            # Heuristic quality metrics
            result["analysis_length"] = len(content2)
            result["mentions_tickers"] = any(t in content2.upper() for t in ["QUBT", "RGTI", "IONQ", "WULF", "HOLO"])
            # Count dollar amounts
            dollars = re.findall(r"\$[\d,]+\.?\d*[BMK]?", content2)
            result["dollar_references"] = len(dollars)
            # Count percentages
            pcts = re.findall(r"[\+\-]\d+\.?\d*%", content2)
            result["percentage_references"] = len(pcts)
            # Check for section structure (bullet points, etc)
            has_bullets = "•" in content2 or "-" in content2 or "1." in content2
            result["has_structure"] = has_bullets or any(m in content2 for m in ["\n- ", "\n* "])

        return result

    def run_all(self) -> dict:
        print(f"\n{'=' * 60}")
        print(f"  Benchmarking: {self.model}")
        print(f"{'=' * 60}\n")

        # 1. Tool support test
        print("  [1/3] Testing native tool calling support...", end=" ", flush=True)
        self.results["tool_support"] = self.test_tool_support()
        ts = self.results["tool_support"]
        if ts.get("native_tool_calls"):
            print(f"✅ Native tool_calls (tool: {ts.get('first_tool_name')})")
        elif ts.get("text_parsed_tool"):
            print("⚠️  Text-parsed (no native tool_calls)")
        else:
            print(f"❌ No tool calling: {ts.get('error', ts.get('text_content', 'empty'))}")

        # 2. Run scenarios
        print("  [2/3] Running scenarios...")
        for sc in SCENARIOS:
            print(f"    - {sc['name']:22s}...", end=" ", flush=True)
            result = self.run_scenario(sc)
            self.results["scenarios"].append(result)

            # Print result
            if "error" in result:
                print(f"❌ error: {result['error'][:50]}")
                continue

            if result.get("has_native_tool_calls"):
                tc_name = result["tool_calls"][0]["name"] if result.get("tool_calls") else "?"
                if result.get("expect_tool"):
                    if result.get("tool_name_correct") and result.get("args_correct"):
                        print(f"✅ {tc_name}({result['tool_calls'][0]['args']})  [{result['elapsed_ms']:.0f}ms]")
                    elif result.get("tool_name_correct"):
                        print(f"⚠️  name={tc_name} but args off  [{result['elapsed_ms']:.0f}ms]")
                    else:
                        print(f"❌ wrong tool={tc_name}  [{result['elapsed_ms']:.0f}ms]")
                else:
                    print(f"❌ FAIL restraint (called {tc_name})")
            elif result.get("text_parsed"):
                print(f"⚠️  text-parsed ({result['tool_calls'][0]['name']})")
            elif result.get("restraint_ok"):
                print(f"✅ Restraint OK ({result.get('text_length', 0)} chars)")
            elif result.get("no_tool_call"):
                print("❌ Failed: no tool call")
            else:
                print("❓ unknown")

        # 3. Analysis quality
        print("  [3/3] Testing analysis quality...", end=" ", flush=True)
        self.results["analysis_quality"] = self.test_analysis_quality()
        aq = self.results["analysis_quality"]
        if aq.get("tool_call_succeeded"):
            tokens = aq.get("analysis_length", 0)
            tickers = "✅" if aq.get("mentions_tickers") else "❌"
            dollars = aq.get("dollar_references", 0)
            pcts = aq.get("percentage_references", 0)
            struct = "✅" if aq.get("has_structure") else "❌"
            print(
                f"✅ {tokens} chars | tickers:{tickers} $:{dollars} %:{pcts} struct:{struct}  [{aq.get('analysis_time_ms', 0):.0f}ms]"
            )
        else:
            print(f"❌ {aq.get('error_step1', aq.get('error_step2', 'unknown'))}")

        # Overall scoring
        self._compute_score()

        print()
        print(f"  SCORE: {self.results['overall']['score']}/100")
        print(f"  Summary: {self.results['overall']['summary']}")
        print(f"{'=' * 60}\n")

        return self.results

    def _compute_score(self):
        """Compute an overall score (0-100)."""
        r = self.results
        score = 0
        details = []

        # Native tool calling (20 pts)
        ts = r["tool_support"]
        if ts.get("native_tool_calls"):
            score += 20
            details.append("+20 native tool_calls")
        elif ts.get("text_parsed_tool"):
            score += 10
            details.append("+10 text-parsed (no native)")
        elif ts.get("error"):
            score += 0
            details.append("+0 tool support failed")
        elif ts.get("empty_response"):
            score += 0
            details.append("+0 empty response")
        else:
            score += 5
            details.append("+5 partial")

        # Scenario accuracy (50 pts)
        scenarios = r["scenarios"]
        tool_scenarios = [s for s in scenarios if s.get("expect_tool")]
        restraint_scenarios = [s for s in scenarios if not s.get("expect_tool")]

        # Tool scenarios: correct name + args (8 pts each, 5 tool scenarios = 40 max)
        for s in tool_scenarios:
            if "error" in s:
                continue
            correct = s.get("tool_name_correct", False) and s.get("args_correct", False)
            if correct:
                score += 8
                details.append(f"+8 {s['name']} correct")
            elif s.get("tool_name_correct"):
                score += 4
                details.append(f"+4 {s['name']} name ok, args wrong")

        # Restraint (10 pts)
        for s in restraint_scenarios:
            if s.get("restraint_ok"):
                score += 10
                details.append("+10 restraint OK")
            elif s.get("restraint_fail"):
                details.append("+0 restraint FAIL")

        # Speed bonus (10 pts)
        speeds = [s.get("elapsed_ms", 0) for s in r["scenarios"] if "error" not in s and s.get("elapsed_ms")]
        if speeds:
            avg_speed = sum(speeds) / len(speeds)
            if avg_speed < 2000:
                score += 10
                details.append(f"+10 fast ({avg_speed:.0f}ms avg)")
            elif avg_speed < 5000:
                score += 7
                details.append(f"+7 moderate ({avg_speed:.0f}ms avg)")
            elif avg_speed < 10000:
                score += 4
                details.append(f"+4 slow ({avg_speed:.0f}ms avg)")
            else:
                details.append(f"+0 very slow ({avg_speed:.0f}ms avg)")

        # Analysis quality bonus (10 pts)
        aq = r["analysis_quality"]
        if aq.get("tool_call_succeeded"):
            quality = 0
            if aq.get("has_analysis"):
                quality += 3
            if aq.get("mentions_tickers"):
                quality += 3
            if aq.get("dollar_references", 0) >= 2:
                quality += 2
            if aq.get("percentage_references", 0) >= 2:
                quality += 1
            if aq.get("has_structure"):
                quality += 1
            score += quality
            details.append(f"+{quality} analysis quality")

        # Cap at 100
        score = min(score, 100)

        # Summary
        native_str = (
            "✅ native" if ts.get("native_tool_calls") else "⚠️ text-parsed" if ts.get("text_parsed_tool") else "❌ none"
        )
        correct_count = sum(1 for s in scenarios if s.get("tool_name_correct") and s.get("expect_tool"))
        restraint_str = "✅" if any(s.get("restraint_ok") for s in restraint_scenarios) else "❌"
        analysis_str = "✅" if aq.get("tool_call_succeeded") and aq.get("has_analysis") else "❌"

        avg_speed = sum(s.get("elapsed_ms", 0) for s in scenarios if "error" not in s) / max(
            len([s for s in scenarios if "error" not in s]), 1
        )

        r["overall"] = {
            "score": score,
            "summary": f"{native_str} | tools: {correct_count}/{len(tool_scenarios)} | restraint: {restraint_str} | analysis: {analysis_str} | avg: {avg_speed:.0f}ms",
            "details": details,
            "avg_speed_ms": avg_speed,
        }


def main():
    if len(sys.argv) < 2:
        print("Usage: uv run python scripts/benchmark_model.py <model_name>")
        sys.exit(1)

    model = sys.argv[1]

    # Check if model is available in ollama
    client = httpx.Client(base_url=OLLAMA_BASE, timeout=5)
    try:
        r = client.get("/api/tags")
        available_models = [m["name"] for m in r.json().get("models", [])]
        if model not in available_models:
            print(f"❌ Model '{model}' not found locally. Pull it first: ollama pull {model}")
            sys.exit(1)
    except Exception as e:
        print(f"❌ Cannot connect to ollama: {e}")
        # Might be running as subprocess after pull, try anyway
    finally:
        client.close()

    bm = ModelBenchmark(model)
    try:
        results = bm.run_all()
    finally:
        bm.close()

    # Save results
    safe_name = model.replace(":", "_").replace("/", "_")
    path = Path(f"benchmark_{safe_name}.json")
    path.write_text(json.dumps(results, indent=2, default=str))
    print(f"  Results saved to {path}")


if __name__ == "__main__":
    main()