"""LLM backend abstraction for benchmark strategies.

Pluggable backends so the benchmark compares orchestration strategies,
not model quality.  All strategies use the same backend instance.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Protocol

logger = logging.getLogger(__name__)


@dataclass
class LLMResponse:
    """Response from an LLM call."""

    text: str
    tokens_in: int = 0
    tokens_out: int = 0
    model: str = ""
    elapsed_seconds: float = 0.0


class LLMBackend(Protocol):
    """Protocol for LLM backends."""

    def generate(
        self, system: str, user: str, temperature: float = 0.2
    ) -> LLMResponse: ...

    @property
    def model_name(self) -> str: ...


class OllamaBackend:
    """Ollama-backed LLM for local benchmarking."""

    def __init__(self, model: str | None = None) -> None:
        from cohort.local.ollama import OllamaClient

        self._client = OllamaClient()
        self._model = model or self._auto_select_model()

    def _auto_select_model(self) -> str:
        models = self._client.list_models()
        # Prefer larger models for better code quality
        for preferred in ["qwen3.5:14b", "qwen3.5:9b", "qwen3.5:2b"]:
            if preferred in models:
                return preferred
        if models:
            return models[0]
        raise RuntimeError("No Ollama models available")

    @property
    def model_name(self) -> str:
        return self._model

    def generate(
        self, system: str, user: str, temperature: float = 0.2
    ) -> LLMResponse:
        result = self._client.generate(
            model=self._model,
            prompt=user,
            system=system,
            temperature=temperature,
            think=True,
        )
        if result is None:
            logger.warning("Ollama returned None for generate call")
            return LLMResponse(text="", model=self._model)
        return LLMResponse(
            text=result.text,
            tokens_in=result.tokens_in,
            tokens_out=result.tokens_out,
            model=result.model,
            elapsed_seconds=result.elapsed_seconds,
        )


class LlamaCppBackend:
    """llama.cpp server backend (OpenAI-compatible API).

    Supports round-robin across multiple server URLs for multi-GPU
    parallelism.  Each generate() call picks the next server.
    """

    def __init__(
        self,
        urls: list[str] | None = None,
        model: str = "default",
    ) -> None:
        import threading

        self._urls = urls or ["http://127.0.0.1:11435", "http://127.0.0.1:11436"]
        self._model = model
        self._index = 0
        self._lock = threading.Lock()

    def _next_url(self) -> str:
        with self._lock:
            url = self._urls[self._index % len(self._urls)]
            self._index += 1
            return url

    @property
    def model_name(self) -> str:
        return self._model

    def generate(
        self, system: str, user: str, temperature: float = 0.2
    ) -> LLMResponse:
        import json
        import time
        import urllib.request

        url = self._next_url()
        t0 = time.monotonic()

        body = json.dumps({
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            "temperature": temperature,
            "max_tokens": 4096,
        }).encode("utf-8")

        req = urllib.request.Request(
            f"{url}/v1/chat/completions",
            data=body,
            headers={"Content-Type": "application/json"},
            method="POST",
        )

        try:
            with urllib.request.urlopen(req, timeout=300) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                elapsed = time.monotonic() - t0
                choice = data.get("choices", [{}])[0]
                text = choice.get("message", {}).get("content", "")
                usage = data.get("usage", {})
                return LLMResponse(
                    text=text,
                    tokens_in=usage.get("prompt_tokens", 0),
                    tokens_out=usage.get("completion_tokens", 0),
                    model=data.get("model", self._model),
                    elapsed_seconds=round(elapsed, 1),
                )
        except Exception as e:
            logger.warning("llama.cpp call to %s failed: %s", url, e)
            return LLMResponse(text="", model=self._model)


class MockBackend:
    """Deterministic mock backend for testing the harness itself."""

    def __init__(self, responses: list[str] | None = None) -> None:
        self._responses = list(responses or [])
        self._call_index = 0

    @property
    def model_name(self) -> str:
        return "mock"

    def generate(
        self, system: str, user: str, temperature: float = 0.2
    ) -> LLMResponse:
        if self._call_index < len(self._responses):
            text = self._responses[self._call_index]
        else:
            text = '```python\nprint("hello")\n```'
        self._call_index += 1
        return LLMResponse(text=text, tokens_in=100, tokens_out=200, model="mock")
