"""Ollama client for the sub-agent, over the one path vmgw exposes.

WHY THE NATIVE /api/chat AND NOT THE OPENAI SHIM

The native endpoint returns `thinking` as its own field rather than
folded into the content, which matters here for two reasons: this model
runs with thinking on and its reasoning is long, and the reasoning must
be kept OUT of the transcript that gets summarised during compaction. A
shim that concatenates them makes both impossible.

THE NUMBERS THAT MATTER, AND WHY

  num_ctx MUST be the resident value (163840). Ollama allocates the KV
  cache per loaded model; asking for a different size makes it RELOAD -
  eighteen gigabytes across the PCIe bus - in the middle of a session.

  keep_alive must match what Open WebUI sends, or this request quietly
  resets the model's idle timer to something shorter and the main
  assistant pays for a reload on its next message.

  The sampling options are sent explicitly rather than inherited. The
  archived sub-agent stack was bitten by exactly this: the Modelfile
  carried presence_penalty 1.5, nothing overrode it, and Ollama merges
  request options OVER Modelfile params - so anything left unsent keeps
  whatever the model row happens to hold that week.

CONTENTION, STATED PLAINLY

OLLAMA_NUM_PARALLEL is 1 and there is one KV slot. Inside this loop that
is cheap: the agent's turns are consecutive, so the prefix cache means
each step prefills only its new tokens. The expensive event is the
handover back to the main model, which finds its cache gone. That is why
the session blocks for as long as it does instead of polling - one long
call costs one cold prefill, ten short ones cost ten.
"""

from __future__ import annotations

import contextlib
import json
import os
import urllib.error
import urllib.request
from typing import Any

BASE_URL = os.environ.get("OWUITOOLS_VMGW_URL", "http://owuitools-vmgw:8000").rstrip("/")
MODEL = os.environ.get("OWUITOOLS_AGENT_MODEL", "qwen3.8:latest")
NUM_CTX = int(os.environ.get("OWUITOOLS_AGENT_NUM_CTX", "163840"))
def _keep_alive(raw: str) -> object:
    """Ollama takes a NUMBER of seconds or a duration STRING, not both.

    `-1` as a number means "keep it loaded indefinitely". The same value
    as the string "-1" is parsed as a duration and rejected:
    `time: missing unit in duration "-1"`. Measured - the first live
    session died on it at step one, having started a VM and done nothing.
    So anything that looks like a plain number is sent as one, and
    anything else ("30m", "1h") is passed through as the string it is.
    """
    text = (raw or "").strip()
    try:
        return int(text)
    except ValueError:
        return text or "5m"


KEEP_ALIVE = _keep_alive(os.environ.get("OWUITOOLS_AGENT_KEEP_ALIVE", "-1"))
REQUEST_TIMEOUT = float(os.environ.get("OWUITOOLS_AGENT_REQUEST_TIMEOUT", "1800"))

# Qwen3's published thinking-mode profile, sent every time so that
# nothing is inherited. See the module docstring.
OPTIONS: dict[str, Any] = {
    "temperature": float(os.environ.get("OWUITOOLS_AGENT_TEMPERATURE", "0.6")),
    "top_p": float(os.environ.get("OWUITOOLS_AGENT_TOP_P", "0.95")),
    "top_k": int(os.environ.get("OWUITOOLS_AGENT_TOP_K", "20")),
    "min_p": 0.0,
    # ZERO, AND NOT BECAUSE THE MAIN MODEL USES ONE.
    #
    # These options were copied from the main assistant's row, where
    # presence_penalty 1.0 is right: it writes prose, and a penalty on
    # tokens already used discourages repetitive phrasing.
    #
    # This agent writes CODE, where the same property is a defect. A
    # 25 KB script says `$form`, `$output`, `Add-Type` and its own
    # function names over and over - that repetition IS the program - and
    # a penalty on it is steady pressure to reach for a different token
    # than the one that is correct. Measured across two failing sessions,
    # both stalled on exactly that shape of bug: "a scope issue with how
    # I'm accessing the log", and "it's using `$a[0]` when `$a` was never
    # actually defined".
    #
    # Qwen3's own guidance says a presence penalty may degrade output and
    # to use it only against repetition loops. The agent already has a
    # repeat guard for those, in tool output, where it can say why.
    "presence_penalty": float(os.environ.get("OWUITOOLS_AGENT_PRESENCE_PENALTY", "0.0")),
    "repeat_last_n": 256,
    "num_ctx": NUM_CTX,
    # A CEILING ON ONE TURN, NOT ON THE SESSION.
    #
    # There was none, so a single turn could run until it chose to stop -
    # measured at twenty-two minutes and roughly fifty thousand tokens,
    # the model designing a whole application inside one reply while the
    # task list sat on item one. This is generous enough that no ordinary
    # turn reaches it and low enough that a runaway one ends. A turn cut
    # short emits no tool call, which the loop already handles by telling
    # the agent to call one.
    "num_predict": int(os.environ.get("OWUITOOLS_AGENT_NUM_PREDICT", "16384")),
}


class OllamaError(RuntimeError):
    pass


def _read_stream(response: Any, on_text: Any) -> dict[str, Any]:
    """Assemble one message out of Ollama's newline-delimited chunks.

    Returns the same shape the non-streaming call returns, so nothing
    downstream can tell the difference. Content and thinking accumulate;
    tool calls are extended as they appear, because Ollama emits a call
    whole rather than a character at a time; the counts arrive on the
    final chunk.

    Anything unreadable in a chunk is skipped rather than raised on. A
    malformed line is one frame of a progress display, and losing the
    whole turn over it would be a poor trade.
    """
    content: list[str] = []
    thinking: list[str] = []
    tool_calls: list[Any] = []
    body: dict[str, Any] = {}
    for raw in response:
        line = raw.decode("utf-8", "replace").strip()
        if not line:
            continue
        try:
            chunk = json.loads(line)
        except ValueError:
            continue
        message = chunk.get("message") or {}
        piece = message.get("content") or ""
        if piece:
            content.append(piece)
            with contextlib.suppress(Exception):
                on_text("".join(content))
        thought = message.get("thinking") or ""
        if thought:
            thinking.append(thought)
            # Only shown when there is nothing else to show: a model that
            # is thinking has not said anything yet, and a blank line for
            # a minute reads as a hung session.
            if not content:
                with contextlib.suppress(Exception):
                    on_text("".join(thinking))
        if message.get("tool_calls"):
            tool_calls.extend(message["tool_calls"])
        if chunk.get("done"):
            body = chunk
    body["message"] = {
        "role": "assistant",
        "content": "".join(content),
        "thinking": "".join(thinking),
        "tool_calls": tool_calls,
    }
    return body


# HOW HARD THE AGENT THINKS, AND WHY IT IS A DIAL RATHER THAN A SWITCH.
#
# Ollama 0.32.15 accepts `think` as a boolean OR one of "high", "medium",
# "low", "max" - its own validator says so in as many words. Whether a
# given model HONOURS a level is a separate question from whether Ollama
# will carry it: levels are native to gpt-oss, and a template that does
# not know about them simply renders the same prompt it always did.
#
# Worth having as a setting either way, because the cost is measurable.
# The agent generates at ~52 tokens/second - fast - and still crawls,
# because a single turn ran to 15,469 tokens, most of it deliberation that
# is then discarded rather than kept. Five minutes of wall clock for one
# step, during which the task list cannot move.
#
# "true" keeps the behaviour that has been solving the hard parts. Set it
# to a level, or to "false", and measure n_gen per turn in Ollama's log
# before believing it did anything.
def _think_setting() -> Any:
    raw = os.environ.get("OWUITOOLS_AGENT_THINK", "true").strip().lower()
    if raw in ("true", "1", "yes"):
        return True
    if raw in ("false", "0", "no"):
        return False
    if raw in ("low", "medium", "high", "max"):
        return raw
    return True


THINK = _think_setting()


def chat(messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
         think: Any = None, on_text: Any = None) -> dict[str, Any]:
    """One turn. Returns the message plus its usage counts.

    STREAMING IS FOR THE PERSON WATCHING, AND FOR NOTHING ELSE.

    This was non-streaming on the grounds that nothing here needs tokens
    as they arrive. That was true of the agent - it acts on a finished
    message - and false of the operator, who asked to see what the model
    is saying while it says it. A turn on this model runs tens of seconds
    with nothing on screen, and "still thinking" is not the same as "still
    working" to somebody deciding whether to wait.

    So `on_text` is the whole of the difference. Without it this streams
    nothing and behaves exactly as before. With it, the reply is read as
    it arrives and the callback is handed the text so far; either way the
    return value is the same assembled message, because the agent is not
    allowed to care which path produced it.

    The callback is the watcher's, not the loop's: it must never be able
    to fail a turn, so it is called inside a suppressed block.
    """
    payload: dict[str, Any] = {
        "model": MODEL,
        "messages": messages,
        "stream": on_text is not None,
        "think": THINK if think is None else think,
        "keep_alive": KEEP_ALIVE,
        "options": OPTIONS,
    }
    if tools:
        payload["tools"] = tools

    request = urllib.request.Request(
        f"{BASE_URL}/ollama/api/chat",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    try:
        with opener.open(request, timeout=REQUEST_TIMEOUT) as response:
            if on_text is None:
                body = json.loads(response.read().decode("utf-8", "replace"))
            else:
                body = _read_stream(response, on_text)
    except urllib.error.HTTPError as exc:
        # Read the body before raising: Ollama puts the real complaint
        # there - an unknown option, a context that will not fit - and
        # losing it leaves only "HTTP 400", which explains nothing.
        detail = exc.read().decode("utf-8", "replace")[:500]
        raise OllamaError(f"the model refused the request (HTTP {exc.code}): {detail}") from exc
    except urllib.error.URLError as exc:
        raise OllamaError(f"could not reach the model: {exc.reason}") from exc

    message = body.get("message") or {}
    return {
        "role": message.get("role", "assistant"),
        "content": message.get("content", "") or "",
        "thinking": message.get("thinking", "") or "",
        "tool_calls": message.get("tool_calls") or [],
        "prompt_tokens": int(body.get("prompt_eval_count") or 0),
        "eval_tokens": int(body.get("eval_count") or 0),
        "prompt_seconds": round((body.get("prompt_eval_duration") or 0) / 1e9, 2),
        "eval_seconds": round((body.get("eval_duration") or 0) / 1e9, 2),
        "done_reason": body.get("done_reason", ""),
    }


def resident_context() -> int | None:
    """What context length the model is actually loaded with.

    Sending a num_ctx that differs from this forces a reload, so the loop
    checks rather than trusting its own default - the two drift the first
    time somebody edits the model row in Open WebUI.
    """
    try:
        opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
        with opener.open(f"{BASE_URL}/ollama/api/ps", timeout=10) as response:
            data = json.loads(response.read().decode())
    except Exception:  # noqa: BLE001
        return None
    for model in data.get("models", []):
        if model.get("name") == MODEL:
            return int(model.get("context_length") or 0) or None
    return None