"""Ollama backend – uses the Ollama REST API (OpenAI-compatible /v1 endpoint)."""

from __future__ import annotations

import copy
import json
from collections.abc import AsyncIterator
from typing import Any, Self

from thryve.core.backends.base import BaseBackend
from thryve.core.backends.ollama.param import OllamaBackendParam
from thryve.core.context.models import (
    BackendInfo,
    ChatResponse,
    FinishReason,
    ImagePart,
    Message,
    TextPart,
    ToolCall,
    UsageInfo,
)
from thryve.utils import get_logger

logger = get_logger("backends.ollama")


class OllamaBackend(BaseBackend):
    """Concrete backend talking to a local Ollama instance via HTTP.

    Ollama exposes an OpenAI-compatible ``/v1/chat/completions`` endpoint, so
    we reuse a thin ``httpx`` layer rather than duplicating the OpenAI SDK.
    """

    def __init__(self, param: OllamaBackendParam) -> None:
        try:
            import httpx  # noqa: F811
        except ImportError as exc:
            raise ImportError("Install 'httpx': pip install httpx") from exc

        self._param = param
        self._base_url = param.base_url.rstrip("/")
        self._client = httpx.Client(base_url=self._base_url, timeout=120.0)
        self._async_client = httpx.AsyncClient(base_url=self._base_url, timeout=120.0)
        self._bound_tools: list[Any] | None = None

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _to_ollama_messages(self, messages: list[Message]) -> list[dict[str, Any]]:
        result: list[dict[str, Any]] = []
        for msg in messages:
            d: dict[str, Any] = {"role": msg.role}
            if isinstance(msg.content, str):
                d["content"] = msg.content
            else:
                parts_text: list[str] = []
                images: list[str] = []
                for part in msg.content:
                    if isinstance(part, TextPart):
                        parts_text.append(part.text)
                    elif isinstance(part, ImagePart):
                        if isinstance(part.data, bytes):
                            images.append(__import__("base64").b64encode(part.data).decode())
                        else:
                            images.append(part.data)
                d["content"] = " ".join(parts_text)
                if images:
                    d["images"] = images

            if msg.tool_calls:
                d["tool_calls"] = [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {"name": tc.name, "arguments": tc.arguments},
                    }
                    for tc in msg.tool_calls
                ]

            if msg.role == "tool" and msg.tool_call_id:
                d["tool_call_id"] = msg.tool_call_id

            result.append(d)
        return result

    def _build_body(self, messages: list[dict[str, Any]], **overrides: Any) -> dict[str, Any]:
        body: dict[str, Any] = {
            "model": self._param.model,
            "messages": messages,
            "stream": False,
            "options": {"temperature": self._param.temperature},
        }
        if self._param.max_tokens is not None:
            body["options"]["num_predict"] = self._param.max_tokens
        if self._bound_tools:
            body["tools"] = [t.to_openai_format() for t in self._bound_tools]
        body.update(overrides)
        return body

    def _parse_response(self, data: dict[str, Any]) -> ChatResponse:
        msg = data.get("message", {})
        content = msg.get("content", "")

        tool_calls: list[ToolCall] | None = None
        raw_tcs = msg.get("tool_calls")
        if raw_tcs:
            tool_calls = []
            for tc in raw_tcs:
                fn = tc.get("function", {})
                args = fn.get("arguments", {})
                if isinstance(args, str):
                    args = json.loads(args)
                tool_calls.append(
                    ToolCall(
                        id=tc.get("id", ""),
                        name=fn.get("name", ""),
                        arguments=args,
                    )
                )

        finish = FinishReason.TOOL_CALLS if tool_calls else FinishReason.STOP

        usage: UsageInfo | None = None
        if "prompt_eval_count" in data:
            usage = UsageInfo(
                prompt_tokens=data.get("prompt_eval_count", 0),
                completion_tokens=data.get("eval_count", 0),
                total_tokens=data.get("prompt_eval_count", 0) + data.get("eval_count", 0),
            )

        return ChatResponse(
            content=content,
            tool_calls=tool_calls,
            backend=self.get_info(),
            finish_reason=finish,
            usage=usage,
        )

    # ------------------------------------------------------------------
    # BaseBackend implementation
    # ------------------------------------------------------------------

    def generate(self, messages: list[Message], **kwargs: Any) -> ChatResponse:
        ollama_msgs = self._to_ollama_messages(messages)
        body = self._build_body(ollama_msgs, **kwargs)
        resp = self._client.post("/api/chat", json=body)
        resp.raise_for_status()
        return self._parse_response(resp.json())

    async def generate_async(self, messages: list[Message], **kwargs: Any) -> ChatResponse:
        ollama_msgs = self._to_ollama_messages(messages)
        body = self._build_body(ollama_msgs, **kwargs)
        resp = await self._async_client.post("/api/chat", json=body)
        resp.raise_for_status()
        return self._parse_response(resp.json())

    async def generate_stream(self, messages: list[Message], **kwargs: Any) -> AsyncIterator[str]:
        if self._bound_tools:
            resp = await self.generate_async(messages, **kwargs)
            if resp.content:
                yield resp.content
            return

        ollama_msgs = self._to_ollama_messages(messages)
        body = self._build_body(ollama_msgs, stream=True)
        async with self._async_client.stream("POST", "/api/chat", json=body) as resp:
            resp.raise_for_status()
            async for line in resp.aiter_lines():
                if not line:
                    continue
                data = json.loads(line)
                chunk = data.get("message", {}).get("content", "")
                if chunk:
                    yield chunk

    def embeddings(self, texts: list[str], **kwargs: Any) -> list[list[float]]:
        results: list[list[float]] = []
        model = kwargs.get("model", self._param.model)
        for text in texts:
            resp = self._client.post("/api/embeddings", json={"model": model, "prompt": text})
            resp.raise_for_status()
            results.append(resp.json().get("embedding", []))
        return results

    def bind_tools(self, tools: list[Any]) -> Self:
        new = copy.copy(self)
        new._bound_tools = list(tools)
        return new  # type: ignore[return-value]

    def get_info(self) -> BackendInfo:
        return BackendInfo(
            provider="ollama",
            model=self._param.model,
            supports_vision=True,
            supports_tools=True,
        )
