"""Benchmark task definitions + graders.

Each task produces a ChatCompletion request and a deterministic grader.
Tasks are intentionally simple and single-turn — this is a compatibility
and cost matrix, not a general capability benchmark.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Any, Callable

from .client import ChatResult


@dataclass
class Task:
    id: str
    description: str
    messages: list[dict[str, Any]]
    grader: Callable[[ChatResult], tuple[bool, str]]
    tools: list[dict[str, Any]] | None = None
    tool_choice: str | dict | None = None
    max_tokens: int = 1024


# ---------- Task 1: reasoning (exact arithmetic) ----------
# Large multiplication — forces the model to reason rather than guess.
# Answer: 31847 * 49123 = 1564490981
def _grade_math(res: ChatResult) -> tuple[bool, str]:
    if not res.ok:
        return False, f"api_error: {res.error}"
    # Accept the answer appearing anywhere in the final content, with or without separators.
    digits = re.sub(r"[^\d]", "", res.content or "")
    expected = "1564420181"
    if expected in digits:
        return True, "exact"
    return False, f"missing_expected (got content: {(res.content or '')[:120]!r})"


task_math = Task(
    id="math_multiplication",
    description="Compute 31847 * 49123 and return the exact integer.",
    messages=[
        {
            "role": "system",
            "content": "You are a precise calculator. Return only the final integer, nothing else.",
        },
        {
            "role": "user",
            "content": "What is 31847 * 49123? Return only the integer.",
        },
    ],
    grader=_grade_math,
    max_tokens=4096,
)


# ---------- Task 2: tool-use (forced single function call) ----------
_WEATHER_TOOL = {
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get the current weather for a city.",
        "parameters": {
            "type": "object",
            "properties": {
                "city": {"type": "string", "description": "City name"},
                "unit": {
                    "type": "string",
                    "enum": ["celsius", "fahrenheit"],
                    "description": "Temperature unit",
                },
            },
            "required": ["city", "unit"],
        },
    },
}


def _grade_tool(res: ChatResult) -> tuple[bool, str]:
    if not res.ok:
        return False, f"api_error: {res.error}"
    if not res.tool_calls:
        return False, f"no_tool_call (content: {(res.content or '')[:120]!r})"
    tc = res.tool_calls[0]
    fn = tc.get("function") or {}
    if fn.get("name") != "get_weather":
        return False, f"wrong_tool: {fn.get('name')!r}"
    try:
        args = json.loads(fn.get("arguments") or "{}")
    except json.JSONDecodeError:
        return False, f"invalid_json_args: {fn.get('arguments')!r}"
    city = (args.get("city") or "").lower()
    unit = (args.get("unit") or "").lower()
    if "paris" not in city:
        return False, f"wrong_city: {city!r}"
    if unit not in {"celsius", "fahrenheit"}:
        return False, f"wrong_unit: {unit!r}"
    return True, f"called get_weather(city={city!r}, unit={unit!r})"


task_tool = Task(
    id="tool_use_weather",
    description="Call get_weather(city='Paris', unit='celsius') via a tool call.",
    messages=[
        {
            "role": "system",
            "content": "You are a weather assistant. Use the provided tool to answer.",
        },
        {
            "role": "user",
            "content": "What is the weather in Paris? Use celsius.",
        },
    ],
    grader=_grade_tool,
    tools=[_WEATHER_TOOL],
    tool_choice="auto",
    max_tokens=512,
)


# ---------- Task 3: coding (valid Python function) ----------
def _grade_code(res: ChatResult) -> tuple[bool, str]:
    if not res.ok:
        return False, f"api_error: {res.error}"
    content = res.content or ""
    # Extract code block if fenced
    m = re.search(r"```(?:python)?\s*(.*?)```", content, re.DOTALL)
    code = m.group(1) if m else content
    # The function must be named is_prime and handle at least these cases.
    try:
        ns: dict[str, Any] = {}
        exec(code, ns)
    except Exception as e:
        return False, f"exec_error: {type(e).__name__}: {e}"
    fn = ns.get("is_prime")
    if not callable(fn):
        return False, "no_function_is_prime"
    cases = [(2, True), (3, True), (4, False), (17, True), (1, False), (0, False), (97, True), (100, False)]
    for n, expected in cases:
        try:
            got = fn(n)
        except Exception as e:
            return False, f"runtime_error on is_prime({n}): {e}"
        if bool(got) != expected:
            return False, f"wrong on is_prime({n}): expected {expected}, got {got}"
    return True, "8/8 cases passed"


task_code = Task(
    id="code_is_prime",
    description="Write a Python function is_prime(n) that returns True/False correctly for 0..100.",
    messages=[
        {
            "role": "system",
            "content": "You write correct, minimal Python. Respond with code only.",
        },
        {
            "role": "user",
            "content": "Write a Python function is_prime(n: int) -> bool. It must handle n <= 1 as False. Return only the code.",
        },
    ],
    grader=_grade_code,
    max_tokens=1024,
)


ALL_TASKS: list[Task] = [task_math, task_tool, task_code]