"""OpenAI-compatible chat client (urllib, no deps) with degenerate-output handling.

Lesson from llmcensor: empty or truncated model output must be a retryable
error, never a silent result. `chat()` retries transport errors and
degenerate completions, then raises DegenerateOutput so the runner records
an errored result instead of a fake score.

Mock models (base_url "mock:") return canned deterministic responses via a
handler registered by the caller — the smoke test registers one that solves
the sample challenge.
"""

from __future__ import annotations

import json
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Callable

from .config import Model


class DegenerateOutput(RuntimeError):
    """Model returned empty/whitespace output or was truncated, after retries."""


class TransportError(RuntimeError):
    """Endpoint unreachable or returned a non-JSON / error response, after retries."""


@dataclass
class Completion:
    content: str
    finish_reason: str
    prompt_tokens: int
    completion_tokens: int


MockHandler = Callable[[Model, list[dict]], str]
_mock_handlers: dict[str, MockHandler] = {}


def register_mock(name: str, handler: MockHandler) -> None:
    _mock_handlers[name] = handler


def chat(model: Model, messages: list[dict], retries: int = 3, timeout: int = 300) -> Completion:
    last_err: Exception | None = None
    for attempt in range(retries):
        if attempt:
            time.sleep(min(2 ** attempt, 30))
        try:
            completion = _request(model, messages, timeout)
        except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, OSError) as err:
            last_err = TransportError(f"{model.alias}: {err}")
            continue
        if not completion.content.strip():
            last_err = DegenerateOutput(
                f"{model.alias}: empty completion (finish_reason={completion.finish_reason})"
            )
            continue
        if completion.finish_reason == "length":
            last_err = DegenerateOutput(
                f"{model.alias}: truncated at max_tokens={model.max_tokens}"
            )
            continue
        return completion
    raise last_err if last_err is not None else TransportError(f"{model.alias}: no attempts made")


def _request(model: Model, messages: list[dict], timeout: int) -> Completion:
    if model.is_mock:
        name = model.base_url.removeprefix("mock:")
        handler = _mock_handlers.get(name)
        if handler is None:
            raise KeyError(f"no mock handler registered for {name!r}")
        return Completion(handler(model, messages), "stop", 0, 0)

    payload = {
        "model": model.model,
        "messages": messages,
        "temperature": model.temperature,
        "max_tokens": model.max_tokens,
    }
    headers = {"Content-Type": "application/json"}
    key = model.api_key()
    if key:
        headers["Authorization"] = f"Bearer {key}"
    req = urllib.request.Request(
        model.base_url.rstrip("/") + "/chat/completions",
        data=json.dumps(payload).encode(),
        headers=headers,
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        body = json.loads(resp.read().decode())
    choice = body["choices"][0]
    usage = body.get("usage", {})
    return Completion(
        content=choice["message"].get("content") or "",
        finish_reason=choice.get("finish_reason") or "stop",
        prompt_tokens=usage.get("prompt_tokens", 0),
        completion_tokens=usage.get("completion_tokens", 0),
    )