"""One OpenAI-compatible provider call, and the model listing beside it.

Chat completions with function tools is the dialect every hosted provider and
every local server speaks — OpenAI, Anthropic's compatibility endpoint,
OpenRouter, Ollama, vLLM — so one client covers them without a provider SDK.
The admin types the URL; the model never chooses where a request goes.

The call is streamed, so the answer reaches the browser as it is written.

Speech is the same dialect's audio half: a realtime transcription session, a
whole clip transcribed, and a sentence spoken, as `parrot serve` answers them.
"""

from __future__ import annotations

import asyncio
import base64
import json
from collections.abc import AsyncIterator
from typing import Any
from urllib.parse import urlparse

import httpx
from websockets.asyncio.client import connect


class ProviderError(Exception):
    """The provider answered, and the answer was a refusal."""


def check_base(api_base: str) -> None:
    parsed = urlparse(api_base)
    if parsed.scheme not in ("http", "https") or not parsed.hostname:
        raise ValueError("The provider URL must be an http(s) URL")


def _headers(api_key: str) -> dict[str, str]:
    return {"Authorization": f"Bearer {api_key}"}


def _detail(response: httpx.Response) -> str:
    try:
        error = response.json().get("error")
    except ValueError:
        error = None
    if isinstance(error, dict):
        error = error.get("message")
    text = str(error or response.text or response.reason_phrase)
    return f"{response.status_code}: {text[:300]}"


class Assembler:
    """The streamed chunks put back together into the message a plain call returns.

    The providers stream the same ``data:`` lines and differ only in the
    margins: OpenRouter interleaves ``:`` comment lines, a thinking model's
    reasoning arrives in a field of its own and is left out, OpenAI and
    Anthropic send a tool call's arguments in pieces where Ollama sends the call
    whole, and some leave out the call's ``index``.
    """

    def __init__(self) -> None:
        self.content = ""
        self.calls: list[dict[str, Any]] = []

    def feed(self, line: str) -> str:
        """One line of the stream; returns the text it adds to the answer."""
        if not line.startswith("data:"):
            return ""
        data = line[5:].strip()
        if not data or data == "[DONE]":
            return ""
        try:
            chunk = json.loads(data)
        except ValueError as exc:
            raise ProviderError(f"sent {data[:300]}") from exc
        error = chunk.get("error")
        if error:
            if isinstance(error, dict):
                error = error.get("message", error)
            raise ProviderError(str(error)[:300])
        choices = chunk.get("choices") or [{}]
        delta = choices[0].get("delta") or {}
        for piece in delta.get("tool_calls") or []:
            self._merge(piece)
        text: str = delta.get("content") or ""
        self.content += text
        return text

    def _merge(self, piece: dict[str, Any]) -> None:
        index = piece.get("index", len(self.calls))
        while len(self.calls) <= index:
            self.calls.append(
                {
                    "id": "",
                    "type": "function",
                    "function": {"name": "", "arguments": ""},
                }
            )
        call = self.calls[index]
        call["id"] = piece.get("id") or call["id"]
        function = piece.get("function") or {}
        call["function"]["name"] = function.get("name") or call["function"]["name"]
        call["function"]["arguments"] += function.get("arguments") or ""

    def message(self) -> dict[str, Any]:
        return {
            "role": "assistant",
            "content": self.content or None,
            "tool_calls": self.calls or None,
        }


async def stream_completion(
    *,
    api_base: str,
    api_key: str,
    model: str,
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
    timeout: float | httpx.Timeout = 120.0,
    extra: dict[str, Any] | None = None,
) -> AsyncIterator[str | dict[str, Any]]:
    """``POST {api_base}/chat/completions`` streamed.

    Yields each piece of the answer's text as it arrives, then the whole
    message. The timeout is per chunk, so a long answer is not cut off.
    ``extra`` is added to the request as it is, for what one provider
    understands and another does not, such as Ollama's ``keep_alive``.
    """
    body: dict[str, Any] = {
        **(extra or {}),
        "model": model,
        "messages": messages,
        "stream": True,
    }
    if tools:
        body["tools"] = tools
    assembler = Assembler()
    async with _client(api_base).stream(
        "POST",
        api_base.rstrip("/") + "/chat/completions",
        json=body,
        headers=_headers(api_key),
        timeout=timeout,
    ) as response:
        if response.is_error:
            await response.aread()
            raise ProviderError(_detail(response))
        async for line in response.aiter_lines():
            text = assembler.feed(line)
            if text:
                yield text
    yield assembler.message()


async def list_models(api_base: str, api_key: str, timeout: float = 15.0) -> list[str]:
    """The ids ``GET {api_base}/models`` offers, for the settings screen."""
    async with httpx.AsyncClient(timeout=timeout) as http:
        response = await http.get(
            api_base.rstrip("/") + "/models", headers=_headers(api_key)
        )
    if response.is_error:
        raise ProviderError(_detail(response))
    data = response.json().get("data") or []
    return sorted(str(m["id"]) for m in data if isinstance(m, dict) and "id" in m)


#: One client per server, so its connection stays open between calls
#: (`parrot serve` keeps an idle one 600 s). A new connection over a lossy
#: link can cost a second, which a spoken reply cannot spare.
_clients: dict[str, httpx.AsyncClient] = {}


def _client(api_base: str) -> httpx.AsyncClient:
    """The client for this server. Used from the engine's loop only."""
    http = _clients.get(api_base)
    if http is None:
        limits = httpx.Limits(keepalive_expiry=300)
        http = _clients[api_base] = httpx.AsyncClient(limits=limits)
    return http


async def warm(api_base: str) -> None:
    """Open the connection to a speech server ahead of the call that needs it."""
    await _client(api_base).get(api_base.rstrip("/") + "/models", timeout=10)


async def transcribe(
    api_base: str,
    audio: bytes,
    filename: str,
    media_type: str,
    *,
    model: str = "whisper-1",
    timeout: float = 30.0,
) -> str:
    """``POST {api_base}/audio/transcriptions``: a whole clip, and what it says."""
    response = await _client(api_base).post(
        api_base.rstrip("/") + "/audio/transcriptions",
        files={"file": (filename, audio, media_type)},
        data={"model": model, "response_format": "json"},
        timeout=timeout,
    )
    if response.is_error:
        raise ProviderError(_detail(response))
    return str(response.json().get("text") or "")


async def speak(
    api_base: str,
    text: str,
    voice: str,
    *,
    model: str = "tts-1",
    timeout: float = 30.0,
) -> bytes:
    """``POST {api_base}/audio/speech``: ``text`` spoken, as MP3.

    MP3 rather than WAV: a seventh of the bytes, which over a slow link to
    the server is most of the wait for a sentence.
    """
    response = await _client(api_base).post(
        api_base.rstrip("/") + "/audio/speech",
        json={"model": model, "input": text, "voice": voice, "response_format": "mp3"},
        timeout=timeout,
    )
    if response.is_error:
        raise ProviderError(_detail(response))
    return response.content


async def transcribe_stream(
    api_base: str,
    audio: asyncio.Queue[bytes | None],
    *,
    rate: int = 16000,
    idle_s: float = 10.0,
) -> AsyncIterator[tuple[bool, str]]:
    """``WS {api_base}/realtime``: one utterance, transcribed while it is spoken.

    Sends each chunk of PCM16 taken off ``audio`` and commits at ``None``.
    Yields ``(False, delta)`` as words are agreed on, then ``(True,
    transcript)`` once. A chunk that takes longer than ``idle_s`` to come
    raises ``TimeoutError``, so an utterance nobody commits ends.
    """
    url = "ws" + api_base.rstrip("/").removeprefix("http") + "/realtime"
    async with connect(url, open_timeout=idle_s) as ws:
        update = {"type": "audio/pcm", "rate": rate}
        await ws.send(
            json.dumps(
                {
                    "type": "session.update",
                    "session": {
                        "type": "transcription",
                        "audio": {"input": {"format": update}},
                    },
                }
            )
        )

        async def append() -> None:
            try:
                while (
                    chunk := await asyncio.wait_for(audio.get(), idle_s)
                ) is not None:
                    data = base64.b64encode(chunk).decode()
                    await ws.send(
                        json.dumps({"type": "input_audio_buffer.append", "audio": data})
                    )
                await ws.send(json.dumps({"type": "input_audio_buffer.commit"}))
            except BaseException:
                # Ends the reading below, which then says why.
                await ws.close()
                raise

        sender = asyncio.create_task(append())
        sender.add_done_callback(lambda done: done.cancelled() or done.exception())
        try:
            async for raw in ws:
                event = json.loads(raw)
                kind = str(event.get("type") or "")
                if kind == "error":
                    raise ProviderError(str((event.get("error") or {}).get("message")))
                if kind.endswith("input_audio_transcription.delta"):
                    yield False, str(event.get("delta") or "")
                elif kind.endswith("input_audio_transcription.completed"):
                    yield True, str(event.get("transcript") or "")
                    return
        finally:
            sender.cancel()
        failure = (
            sender.exception() if sender.done() and not sender.cancelled() else None
        )
        raise failure or ProviderError("the session closed before its transcript")