from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import httpx

from harness.adapters import Adapter
from harness.verifier import verify

PROMPTS = Path(__file__).resolve().parent / "prompts"
# The arena gives 120 s per request. Long inputs may take up to 75 s end to end,
# so the model call gets 100 s and the rest is left for our own work.
SLM_TIMEOUT_S = 100.0


@dataclass
class HarnessResult:
    answer: Any
    verified: bool
    errors: str = ""
    trace: str = ""
    steps: int = 0
    attempts: list[str] = field(default_factory=list)


def complete(
    slm_url: str, prompt: str, max_tokens: int, stop: list[str], seed: int | None = None
) -> tuple[str, str]:
    """One completion call. Returns (text, finish_reason); raises on transport errors."""
    body: dict[str, Any] = {
        "model": "slm",
        "prompt": prompt,
        "max_tokens": max_tokens,
        "stop": stop,
        "temperature": 0,
    }
    if seed is not None:
        body["seed"] = seed
    with httpx.Client(timeout=SLM_TIMEOUT_S) as http:
        response = http.post(slm_url.rstrip("/") + "/v1/completions", json=body)
    response.raise_for_status()
    choice = (response.json().get("choices") or [{}])[0]
    return str(choice.get("text") or ""), str(choice.get("finish_reason") or "")


def _run_compression(
    task: dict[str, Any], *, adapter: Adapter, slm_url: str, config: dict[str, Any], seed: int | None
) -> HarnessResult:
    """ctx_compress: exactly one full model pass, then deterministic repair.

    The arena allows two model calls per item (one full pass, one repair pass
    on the output only). Deterministic repair is faster, cheaper and never
    invents anything, so the second call is not used. About 1 in 7 model
    responses are faulty on purpose; any failure falls through to a summary
    built from the session itself instead of crashing.
    """
    max_tokens = int((config.get("slm") or {}).get("max_new_tokens") or 3072)
    prefill = getattr(adapter, "prefill", "")
    trace: list[str] = []
    draft = ""
    try:
        prompt = adapter.render(PROMPTS, "task", task)
        text, reason = complete(slm_url, prompt, max_tokens, adapter.stop, seed)
        draft = prefill + text if prefill and not text.lstrip().startswith("## ") else text
        trace.append(f"call 1: {len(text)} chars, finish {reason or 'unknown'}")
    except Exception as exc:  # timeouts, 5xx, bad JSON: never crash
        trace.append(f"call 1 failed: {type(exc).__name__}")
    try:
        answer, note = adapter.finish(draft, task)  # type: ignore[attr-defined]
    except Exception as exc:
        trace.append(f"finish failed: {type(exc).__name__}")
        answer, note = adapter.finish("", task)  # type: ignore[attr-defined]
    trace.append(f"repair: {note or 'none needed'}")
    verified = bool(draft.strip()) and "built from the session" not in note
    return HarnessResult(
        answer=answer,
        verified=verified,
        errors=note,
        trace="\n".join(trace),
        steps=1,
        attempts=[draft],
    )


def run_harness(
    task: dict[str, Any],
    *,
    adapter: Adapter,
    slm_url: str,
    tools_url: str,
    config: dict[str, Any],
    seed: int | None = None,
) -> HarnessResult:
    if hasattr(adapter, "finish"):
        return _run_compression(task, adapter=adapter, slm_url=slm_url, config=config, seed=seed)

    # other arenas: the template's verify and retry loop
    max_tokens = int((config.get("slm") or {}).get("max_new_tokens") or 256)
    max_steps = int((config.get("harness") or {}).get("max_steps") or 3)
    errors = ""
    answer = ""
    trace: list[str] = []
    attempts: list[str] = []
    for step in range(1, max_steps + 1):
        kind = "task" if not errors else "retry"
        prompt = adapter.render(PROMPTS, kind, task, draft=answer, errors=errors, tools_url=tools_url)
        try:
            answer, _ = complete(slm_url, prompt, max_tokens, adapter.stop, seed)
        except Exception as exc:
            trace.append(f"step {step}: call failed {type(exc).__name__}")
            errors = "the model call failed"
            continue
        attempts.append(answer)
        verdict = verify(adapter, answer, task, tools_url)
        trace.append(f"step {step}: {'ok' if verdict.ok else verdict.reason}")
        if verdict.ok:
            return HarnessResult(verdict.answer, True, trace="\n".join(trace), steps=step, attempts=attempts)
        errors = verdict.reason
    return HarnessResult(adapter.clean(answer), False, errors, "\n".join(trace), max_steps, attempts)