# SPDX-License-Identifier: GPL-2.0-only
#
# llamachat - a small native chat client for a local llama.cpp router
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
"""Talking to the llama.cpp router over its OpenAI-compatible API."""

import base64
import hashlib
import io
import json
import mimetypes
from dataclasses import dataclass, field
from pathlib import Path

import httpx
from PIL import Image

from . import search
from . import providers as providers_mod

THUMB_SIZE = (128, 128)

# Extensions we are willing to read as text even though mimetypes does not
# call them text/*.
CODE_SUFFIXES = {
    ".py", ".c", ".h", ".cpp", ".hpp", ".rs", ".go", ".js", ".ts", ".jsx",
    ".tsx", ".sh", ".bash", ".zsh", ".lua", ".rb", ".pl", ".php", ".java",
    ".kt", ".swift", ".sql", ".toml", ".yaml", ".yml", ".json", ".ini",
    ".cfg", ".conf", ".md", ".rst", ".txt", ".csv", ".xml", ".html", ".css",
    ".diff", ".patch", ".SlackBuild", ".info", ".desktop", ".service",
}

IMAGE_MIMES = {"image/png", "image/jpeg", "image/gif", "image/webp", "image/bmp"}

# Sentinel yielded by the SSE parser at the end of a stream. A distinct
# object rather than '' so an empty delta cannot be mistaken for the end.
DONE = ("done", "")


class BackendError(Exception):
    """Any failure talking to the router, already phrased for the user."""


@dataclass
class SearchConfig:
    """What stream_chat needs to know to offer and run web search."""
    enabled: bool = False
    url: str = ""
    results: int = 5
    snippet_chars: int = 300
    timeout: int = 10
    max_searches: int = 2


@dataclass
class Attachment:
    """A file the user attached, ready for both the API and the database."""
    path: Path
    kind: str            # 'text' | 'image'
    mime: str
    size: int
    sha256: str
    text: str = ""       # kind == 'text': file content, possibly truncated
    data_url: str = ""   # kind == 'image': base64 data: URL for the API
    thumb: bytes | None = None
    truncated: bool = False


def classify(path: Path) -> str:
    """Decide whether a path is a text or image attachment."""
    mime, _ = mimetypes.guess_type(path.name)
    if mime in IMAGE_MIMES:
        return "image"
    if path.suffix in CODE_SUFFIXES:
        return "text"
    if mime and mime.startswith("text/"):
        return "text"
    return "unknown"


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def load_attachment(path: Path, char_budget: int) -> Attachment:
    """Read a file into an Attachment, truncating text past char_budget."""
    path = path.expanduser().resolve()
    kind = classify(path)
    if kind == "unknown":
        raise BackendError(f"Unsupported file type: {path.name}")

    size = path.stat().st_size
    mime = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
    att = Attachment(
        path=path, kind=kind, mime=mime, size=size, sha256=_sha256(path)
    )

    if kind == "text":
        raw = path.read_text(encoding="utf-8", errors="replace")
        if len(raw) > char_budget:
            att.text = raw[:char_budget]
            att.truncated = True
        else:
            att.text = raw
        return att

    raw = path.read_bytes()
    att.data_url = f"data:{mime};base64,{base64.b64encode(raw).decode('ascii')}"
    att.thumb = _thumbnail(path)
    return att


def _thumbnail(path: Path) -> bytes | None:
    """A small JPEG preview, so reopened chats still show what was sent."""
    try:
        with Image.open(path) as img:
            img = img.convert("RGB")
            img.thumbnail(THUMB_SIZE)
            buf = io.BytesIO()
            img.save(buf, format="JPEG", quality=80)
            return buf.getvalue()
    except Exception:
        # ponytail: a missing preview is cosmetic, never fail an attach over it
        return None


def build_user_content(text: str, attachments: list[Attachment]):
    """Assemble one OpenAI-format user message from text plus attachments.

    Returns a plain string when there is nothing but text, and the
    multi-part content array when images are involved.
    """
    parts: list[str] = []
    for att in attachments:
        if att.kind != "text":
            continue
        note = " (truncated)" if att.truncated else ""
        parts.append(
            f"--- file: {att.path.name}{note} ---\n{att.text}\n--- end file ---"
        )
    if text:
        parts.append(text)
    joined = "\n\n".join(parts)

    images = [a for a in attachments if a.kind == "image"]
    if not images:
        return joined

    content: list[dict] = [{"type": "text", "text": joined}]
    for att in images:
        content.append(
            {"type": "image_url", "image_url": {"url": att.data_url}}
        )
    return content


class Client:
    """Minimal OpenAI-compatible client for the local router."""

    def __init__(self, base_url: str, timeout: int = 300, api_key: str = ""):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.api_key = api_key

    def _headers(self) -> dict:
        """Bearer auth when the provider needs it, nothing when it does not."""
        return {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}

    def models(self, list_timeout: int = 30) -> list[str]:
        """Model ids the router currently offers.

        `list_timeout` is the overall deadline for the listing call, separate
        from the per-turn `self.timeout` used by chat: the picker has to
        populate fast, and a dead provider must not hang it for minutes.
        The connect deadline is shorter still, so an unreachable host is
        reported quickly while a reachable one still gets the full window.
        """
        try:
            resp = httpx.get(
                f"{self.base_url}/v1/models",
                timeout=httpx.Timeout(list_timeout, connect=5),
                headers=self._headers(),
            )
            resp.raise_for_status()
            payload = resp.json()
        except httpx.HTTPError as exc:
            raise BackendError(f"Cannot reach router at {self.base_url}: {exc}")
        except ValueError as exc:
            raise BackendError(f"Router sent invalid JSON: {exc}")
        return [m["id"] for m in payload.get("data", []) if "id" in m]

    def complete(self, model: str, messages: list[dict], max_tokens: int = 48) -> str:
        """One short non-streaming reply, for side errands like titling.

        No tools are offered and failures are the caller's to catch: this is
        never the user's own turn, so a dead router must not break the chat.

        Thinking is turned off through the chat template: a reasoning model
        otherwise spends the whole budget on its thoughts and returns an
        empty content field. `reasoning_budget: 0` does not do this, it ends
        the thinking immediately and still yields nothing.
        """
        try:
            resp = httpx.post(
                f"{self.base_url}/v1/chat/completions",
                json={
                    "model": model,
                    "messages": messages,
                    "stream": False,
                    "max_tokens": max_tokens,
                    "chat_template_kwargs": {"enable_thinking": False},
                },
                timeout=httpx.Timeout(self.timeout, connect=10),
                headers=self._headers(),
            )
            resp.raise_for_status()
            payload = resp.json()
        except httpx.HTTPError as exc:
            raise BackendError(f"Request failed: {exc}")
        except ValueError as exc:
            raise BackendError(f"Router sent invalid JSON: {exc}")
        choices = payload.get("choices") or []
        if not choices:
            return ""
        return str((choices[0].get("message") or {}).get("content") or "")

    def stream_chat(
        self, model: str, messages: list[dict], search_cfg: "SearchConfig | None" = None
    ):
        """Yield (kind, text) pairs as the model produces them.

        `kind` is 'reasoning' for the model's thinking and 'content' for the
        reply proper. Presets with a reasoning budget emit the former in a
        separate `reasoning_content` delta field, which is what lets the UI
        keep the two apart.

        With `search_cfg`, the model is offered a `web_search` tool and this
        becomes a loop: a round that ends in a tool call is searched for, the
        result appended, and the request re-sent. Two further kinds are
        yielded around each search, 'search_start' and 'search_done'. After
        `max_searches` rounds the tool is withdrawn, which forces an answer.

        Loading a different model makes the router unload the previous one,
        so the first chunk can take several seconds. That wait happens
        inside the initial `stream()` call.
        """
        if search_cfg is None or not search_cfg.enabled:
            yield from self._stream_once(model, messages, tools=None)
            return

        # The loop mutates its own copy; the caller's list is history.
        convo = list(messages)
        for round_number in range(search_cfg.max_searches + 1):
            last_round = round_number == search_cfg.max_searches
            tools = None if last_round else [search.TOOL_SCHEMA]
            calls: dict = {}
            saw_tool_finish = False

            for kind, piece in self._stream_once(model, convo, tools=tools):
                if kind == "tool_calls":
                    search.accumulate(calls, json.loads(piece))
                elif kind == "tool_finish":
                    saw_tool_finish = True
                else:
                    yield (kind, piece)

            wanted = [
                call
                for call in calls.values()
                if call["name"] == search.TOOL_SCHEMA["function"]["name"]
            ]
            if not (saw_tool_finish and wanted) or last_round:
                return

            convo.append(
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": call["id"],
                            "type": "function",
                            "function": {
                                "name": call["name"],
                                "arguments": call["arguments"],
                            },
                        }
                        for call in wanted
                    ],
                }
            )
            # The round after this one is the last, so its results are the
            # last the model will get. Saying so is what makes it answer
            # instead of asking for a search it can no longer make.
            final = round_number + 1 == search_cfg.max_searches
            for call in wanted:
                yield from self._run_search(convo, call, search_cfg, final)

    def _run_search(
        self, convo: list[dict], call: dict, cfg: "SearchConfig", last: bool = False
    ):
        """Perform one search, append its result to `convo`, report both ends.

        Every failure still appends a tool message: the model is waiting on
        a result it asked for, and the turn has to complete either way.
        """
        query = search.parse_query(call["arguments"])
        if not query:
            # Unusable arguments: nothing to search, but the call still needs
            # an answer or the next request is malformed.
            convo.append(
                search.tool_message(
                    call["id"], None, error="Malformed search request", last=last
                )
            )
            return

        yield ("search_start", query)
        try:
            results = search.search(
                cfg.url, query, cfg.results, cfg.snippet_chars, cfg.timeout
            )
        except search.SearchError as exc:
            convo.append(
                search.tool_message(call["id"], None, error=str(exc), last=last)
            )
            yield (
                "search_done",
                json.dumps({"query": query, "results": [], "error": str(exc)}),
            )
            return
        convo.append(search.tool_message(call["id"], results, last=last))
        yield (
            "search_done",
            json.dumps({"query": query, "results": results, "error": ""}),
        )

    def _stream_once(self, model: str, messages: list[dict], tools: list | None):
        """One request/response round, yielding parsed SSE pairs."""
        body = {
            "model": model,
            "messages": messages,
            "stream": True,
            # The final chunk then carries exact prompt/completion counts,
            # which is what makes the context meter truthful rather than an
            # estimate, at no extra request.
            "stream_options": {"include_usage": True},
        }
        if tools:
            body["tools"] = tools
        try:
            with httpx.stream(
                "POST",
                f"{self.base_url}/v1/chat/completions",
                json=body,
                timeout=httpx.Timeout(self.timeout, connect=10),
                headers=self._headers(),
            ) as resp:
                if resp.status_code != 200:
                    resp.read()
                    raise BackendError(
                        f"Router returned {resp.status_code}: {resp.text[:300]}"
                    )
                for line in resp.iter_lines():
                    chunk = _parse_sse_line(line)
                    if chunk is None:
                        continue
                    if chunk == DONE:
                        return
                    yield chunk
        except httpx.HTTPError as exc:
            raise BackendError(f"Request failed: {exc}")


class MultiClient:
    """One facade over every configured provider.

    Holds a `Client` per provider, built on demand so a key is resolved only
    when that provider is actually used. The UI talks in prefixed model ids
    and never needs to know which endpoint one lives on.
    """

    def __init__(
        self, table, timeout=300, list_timeout=30, resolver=None,
        client_factory=Client,
    ):
        self.table = table
        self.timeout = timeout
        self.list_timeout = list_timeout
        self.resolver = resolver or providers_mod.KeyResolver()
        self._factory = client_factory
        self._clients: dict[str, Client] = {}

    def client_for(self, model_id: str) -> Client:
        """The client that serves this model, resolving its key on first use."""
        name, _ = providers_mod.split(model_id, self.table)
        if name not in self._clients:
            provider = self.table[name]
            self._clients[name] = self._factory(
                provider.base_url,
                timeout=self.timeout,
                api_key=self.resolver.resolve(provider),
            )
        return self._clients[name]

    def wire_name(self, model_id: str) -> str:
        """The model name the provider itself expects, without our prefix."""
        _, model = providers_mod.split(model_id, self.table)
        return model

    def models(self) -> tuple[list[str], list[str]]:
        """Every offered model id, plus notes about what went wrong.

        A provider that is unreachable or whose filter matched nothing must
        not stop the others being listed: local models have to stay usable
        when the network is down.

        The fan-out is serial on purpose: `KeyResolver._cache` is unlocked
        and only safe while every caller resolves on the GUI thread. Threading
        this would spawn two pinentry prompts for one hardware token.
        """
        listed: list[str] = []
        problems: list[str] = []
        for name, provider in self.table.items():
            try:
                available = self._listing_client(provider).models(
                    list_timeout=self.list_timeout
                )
            except (BackendError, providers_mod.KeyResolutionError) as exc:
                problems.append(f"{name}: {exc}")
                continue
            kept = providers_mod.apply_filter(provider, available)
            if available and not kept:
                problems.append(
                    f"{name}: 0 of {len(available)} models matched filter"
                )
            listed.extend(providers_mod.qualify(name, m) for m in kept)
        return listed, problems

    def _listing_client(self, provider) -> Client:
        """Listing needs a client too, and needs the key for a private API."""
        return self.client_for(providers_mod.qualify(provider.name, ""))


# How much of the first exchange the titling request is shown. Enough to
# see what the conversation is about, far short of a whole context.
TITLE_EXCERPT = 1200
TITLE_MAX_CHARS = 60

_TITLE_SYSTEM = (
    "You write short titles for chat conversations. Reply with the title "
    "alone: at most six words, no quotes, no punctuation at the end, no "
    "explanation, no preamble."
)


def placeholder_title(text: str) -> str:
    """The stand-in title a new session gets: the question's opening words."""
    return text[:TITLE_MAX_CHARS]


def needs_title(current_title: str, first_question: str, reply_text: str) -> bool:
    """Whether a finished turn should be handed to the titler.

    Keyed on the session still wearing the placeholder derived from its
    opening question, rather than on being the first turn. A first reply can
    come back empty, and a one-shot window appends every question to the
    same session until New is pressed, so "first turn" gives up too early in
    both cases.
    """
    if not reply_text.strip():
        return False
    return current_title == placeholder_title(first_question)


def title_messages(user_text: str, reply_text: str) -> list[dict]:
    """The request that asks the model to title its own first exchange."""
    exchange = (
        f"User: {user_text[:TITLE_EXCERPT]}\n\n"
        f"Assistant: {reply_text[:TITLE_EXCERPT]}"
    )
    return [
        {"role": "system", "content": _TITLE_SYSTEM},
        {
            "role": "user",
            "content": (
                f"{exchange}\n\nTitle this conversation in at most six words."
            ),
        },
    ]


def clean_title(raw: str) -> str:
    """Reduce a model's reply to a usable title, or '' if there is none.

    Models decorate: they think out loud, add a "Title:" label, wrap the
    answer in quotes or bold, or chat before answering. Everything here is
    salvage, and returning '' is a valid outcome that leaves the caller's
    fallback title in place.
    """
    text = raw or ""
    if "<think>" in text:
        # An unclosed block means the reply was cut off mid-thought and
        # nothing after it can be trusted as a title.
        _, _, after = text.partition("</think>")
        text = after if "</think>" in raw else ""

    # Chatter comes before the answer, so the last non-empty line is it.
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    if not lines:
        return ""
    title = lines[-1]

    title = title.strip("*_` ")
    for label in ("title:", "Title:", "TITLE:"):
        if title.startswith(label):
            title = title[len(label):].strip()
    title = title.strip().strip('"“”\'')
    title = title.rstrip(".!,;:").strip()

    if len(title) > TITLE_MAX_CHARS:
        cut = title[:TITLE_MAX_CHARS]
        # Prefer a word boundary, but never return a single truncated word.
        head, sep, _ = cut.rpartition(" ")
        title = (head if sep and head else cut).rstrip(".,;:-")
    return title.strip()


def estimate_tokens(messages: list[dict], chars_per_token: float) -> int:
    """Rough token count for a request that has not been sent yet.

    Used for the live meter while typing. Once a reply completes, the exact
    count from the usage block replaces this. Images are counted as a flat
    allowance rather than by dimension, which is enough to show that a
    picture is expensive without pretending to know the tiling.
    """
    chars = 0
    images = 0
    for message in messages:
        chars += len(str(message.get("role", "")))
        content = message.get("content")
        if isinstance(content, str):
            chars += len(content)
            continue
        for part in content or []:
            if part.get("type") == "text":
                chars += len(part.get("text", ""))
            else:
                images += 1
    # A few tokens per message go to the chat template's own markup.
    overhead = 4 * len(messages)
    return int(chars / max(chars_per_token, 1.0)) + overhead + images * 600


def _parse_sse_line(line: str) -> tuple[str, str] | None:
    """Decode one SSE line.

    Returns a ('reasoning'|'content', text) pair, the DONE sentinel at the
    end of the stream, or None for lines carrying nothing to render.
    """
    if not line.startswith("data:"):
        return None
    payload = line[5:].strip()
    if payload == "[DONE]":
        return DONE
    try:
        obj = json.loads(payload)
    except ValueError:
        return None
    # The usage chunk arrives with an empty choices list, so check it first.
    usage = obj.get("usage")
    if usage and usage.get("prompt_tokens") is not None:
        return ("usage", json.dumps(usage))

    choices = obj.get("choices") or []
    if not choices:
        return None
    delta = choices[0].get("delta") or {}

    # Tool-call fragments arrive in their own delta field, alongside a
    # finish_reason of 'tool_calls' on the final chunk of the round.
    tool_calls = delta.get("tool_calls")
    if tool_calls:
        return ("tool_calls", json.dumps(tool_calls))
    if choices[0].get("finish_reason") == "tool_calls":
        return ("tool_finish", "")

    # Both fields can arrive in one delta, and a reasoning delta can be an
    # empty string. Test for presence, not truthiness, so a chunk carrying
    # reasoning is never mistaken for a content chunk.
    reasoning = delta.get("reasoning_content")
    if reasoning is not None:
        return ("reasoning", reasoning) if reasoning else None
    content = delta.get("content")
    if content:
        return ("content", content)
    return None
