"""Model completions via Ollama's raw /api/generate endpoint.

We deliberately do *not* go through the agent's chat provider here: the
capability benchmark wants a raw HumanEval-style completion of the function
prefix (not a chat turn), and we want full control of the read timeout — the
agent's OllamaProvider hard-codes 120s, which throttles reasoning models. This
module owns its own timeout knob.
"""

from __future__ import annotations

from dataclasses import dataclass

import requests


@dataclass
class Completion:
    text: str
    ok: bool
    error: str | None = None
    elapsed: float = 0.0


def complete(model: str, prompt: str, stop: list[str] | None = None,
             *, host: str = "http://127.0.0.1:11434", temperature: float = 0.2,
             num_predict: int = 512, timeout: float = 300.0,
             seed: int | None = None) -> Completion:
    """Ask the model to continue `prompt`. Stop tokens are passed to Ollama and
    re-applied client-side (Ollama strips the stop string, which is what we want
    — the assembled program must not contain the test's leading token twice).

    ``seed`` is forwarded to Ollama's sampler so multi-seed runs at the same
    temperature draw *different but reproducible* samples (the basis for pass@1
    confidence intervals)."""
    import time
    t0 = time.time()
    options = {"temperature": temperature, "num_predict": num_predict}
    if stop:
        options["stop"] = stop
    if seed is not None:
        options["seed"] = seed
    try:
        # raw=True bypasses the chat template so an instruct model *continues*
        # the code (HumanEval-style) instead of replying conversationally with
        # prose + markdown fences, which is what MultiPL-E's assembly expects.
        r = requests.post(f"{host}/api/generate", json={
            "model": model, "prompt": prompt, "stream": False,
            "options": options, "raw": True}, timeout=timeout)
        r.raise_for_status()
        text = r.json().get("response", "")
    except Exception as e:  # noqa: BLE001 — surface as a failed completion row
        return Completion("", False, str(e), time.time() - t0)
    return Completion(_truncate(text, stop), True, None, time.time() - t0)


def _strip_fences(text: str) -> str:
    """Drop a markdown code fence even when raw=True fails to suppress it.

    Chatty coder models sometimes wrap the continuation in ``` fences (and append
    a prose ``# Explanation`` after a closing fence). A triple-backtick is not
    valid syntax in any target language, so it is a safe, unambiguous cut point.
    Two shapes occur: a *wrapped* block (opens with ```lang … closes with ```)
    and a *bare* continuation that the model terminates with a stray closing ```.
    """
    fence = "```"
    if fence not in text:
        return text
    stripped = text.lstrip()
    if stripped.startswith(fence):
        # wrapped: drop the opening ```lang line, keep until the next fence.
        body = stripped[len(fence):]
        nl = body.find("\n")
        body = body[nl + 1:] if nl != -1 else ""
        end = body.find(fence)
        return body[:end] if end != -1 else body
    # bare continuation: cut at the stray closing fence.
    return text[:text.find(fence)]


def _truncate(text: str, stop: list[str] | None) -> str:
    """Defensive client-side stop truncation (covers the no-stop / streamed
    cases and any model that ignores the option)."""
    text = _strip_fences(text)
    if not stop:
        return text
    cut = len(text)
    for s in stop:
        i = text.find(s)
        if i != -1:
            cut = min(cut, i)
    return text[:cut]