"""Local Ollama Provider 适配器.

Ollama 本地部署（http://localhost:11434），无需 API key。
调用 /api/chat 端点，支持流式与非流式。
"""

import json
from collections.abc import AsyncGenerator
from typing import Any

import httpx
import structlog

from ..errors import AILLMUnavailableError
from .base import LLMProvider, LLMResponse, LLMStreamChunk

logger = structlog.get_logger()


class LocalOllamaProvider(LLMProvider):
    """本地 Ollama Provider.

    Ollama API 格式与 OpenAI 略有不同：
      - 端点：/api/chat（非 /v1/chat/completions）
      - 消息格式：{"role": "...", "content": "..."}（兼容）
      - 流式：逐行 JSON（非 SSE）
    """

    def __init__(
        self,
        base_url: str = "http://localhost:11434",
        timeout: float = 60.0,
        stream_connect_timeout: float = 30.0,
        stream_read_timeout: float = 120.0,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._stream_connect_timeout = stream_connect_timeout
        self._stream_read_timeout = stream_read_timeout

    @property
    def name(self) -> str:
        return "local_ollama"

    def is_available(self) -> bool:
        return bool(self._base_url)

    def _build_url(self) -> str:
        return f"{self._base_url}/api/chat"

    async def chat(
        self,
        messages: list[dict[str, str]],
        model: str,
        temperature: float = 0.7,
        **kwargs: Any,
    ) -> LLMResponse:
        if not self.is_available():
            raise AILLMUnavailableError("ollama base_url not configured")

        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "options": {"temperature": temperature},
            "stream": False,
        }
        payload.update(kwargs)

        try:
            async with httpx.AsyncClient(timeout=self._timeout) as client:
                resp = await client.post(self._build_url(), json=payload)
                resp.raise_for_status()
                data = resp.json()
        except httpx.HTTPStatusError as exc:
            logger.error(
                "ollama_chat_http_error",
                status_code=exc.response.status_code,
                body=exc.response.text[:500],
            )
            raise AILLMUnavailableError(
                f"ollama http {exc.response.status_code}: {exc.response.text[:200]}",
            ) from exc
        except httpx.HTTPError as exc:
            logger.error("ollama_chat_http_error", error=str(exc))
            raise AILLMUnavailableError(f"ollama http error: {exc}") from exc

        message = data.get("message", {})
        content = message.get("content", "") or ""
        # Ollama 返回 eval_count 等，转换为通用 usage
        prompt_tokens = int(data.get("prompt_eval_count", 0))
        completion_tokens = int(data.get("eval_count", 0))
        usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        }
        return LLMResponse(
            content=content,
            model=data.get("model", model),
            usage=usage,
            provider=self.name,
            raw=data,
        )

    async def stream_chat(
        self,
        messages: list[dict[str, str]],
        model: str,
        temperature: float = 0.7,
        **kwargs: Any,
    ) -> AsyncGenerator[LLMStreamChunk, None]:
        if not self.is_available():
            raise AILLMUnavailableError("ollama base_url not configured")

        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "options": {"temperature": temperature},
            "stream": True,
        }
        payload.update(kwargs)

        timeout = httpx.Timeout(
            connect=self._stream_connect_timeout,
            read=self._stream_read_timeout,
            write=self._stream_connect_timeout,
            pool=self._stream_connect_timeout,
        )

        try:
            async with (
                httpx.AsyncClient(timeout=timeout) as client,
                client.stream("POST", self._build_url(), json=payload) as resp,
            ):
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    chunk = self._parse_ndjson_line(line, model)
                    if chunk is not None:
                        yield chunk
                        if chunk.finish_reason == "stop":
                            return
        except httpx.HTTPStatusError as exc:
            logger.error(
                "ollama_stream_http_error",
                status_code=exc.response.status_code,
            )
            raise AILLMUnavailableError(
                f"ollama stream http {exc.response.status_code}",
            ) from exc
        except httpx.HTTPError as exc:
            logger.error("ollama_stream_http_error", error=str(exc))
            raise AILLMUnavailableError(f"ollama stream error: {exc}") from exc

    @staticmethod
    def _parse_ndjson_line(line: str, model: str) -> LLMStreamChunk | None:
        """解析 Ollama NDJSON 格式（每行一个 JSON 对象）.

        Ollama 流式响应格式：
            {"model":"...","message":{"role":"assistant","content":"..."},"done":false}
            {"model":"...","message":{"role":"assistant","content":""},"done":true}
        """
        if not line:
            return None
        try:
            data = json.loads(line)
        except json.JSONDecodeError:
            return None
        message = data.get("message", {})
        delta = message.get("content", "") or ""
        done = data.get("done", False)
        return LLMStreamChunk(
            delta=delta,
            model=data.get("model", model),
            finish_reason="stop" if done else None,
            provider="local_ollama",
        )

    async def embed(self, text: str, model: str = "nomic-embed-text") -> list[float]:
        """文本向量化（Ollama embedding API）."""
        if not self.is_available():
            raise AILLMUnavailableError("ollama base_url not configured")
        url = f"{self._base_url}/api/embeddings"
        try:
            async with httpx.AsyncClient(timeout=self._timeout) as client:
                resp = await client.post(url, json={"model": model, "prompt": text})
                resp.raise_for_status()
                data = resp.json()
        except httpx.HTTPError as exc:
            logger.error("ollama_embed_error", error=str(exc))
            raise AILLMUnavailableError(f"ollama embed error: {exc}") from exc
        return data.get("embedding", [])
