from __future__ import annotations

import json
from typing import Any, Iterator

from niannian.llm.providers.base import default_stream_complete, http_json_request, http_stream_request
from niannian.llm.types import CompletionResult, EmbeddingResult, Message, OcrResult, RerankResult, StreamChunk, Usage


class OpenAICompatProvider:
    """OpenAI / DeepSeek / 通义 / 月之暗面 等兼容 API。"""

    name = "openai_compat"

    def __init__(self, config: dict[str, Any]) -> None:
        self._api_key = (config.get("api_key") or "").strip()
        base = config.get("base_url") or "https://api.openai.com/v1"
        self._base_url = str(base).rstrip("/")

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

    def _headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }

    def complete(
        self,
        model: str,
        messages: list[Message],
        *,
        temperature: float = 0.7,
        max_tokens: int | None = None,
        json_mode: bool = False,
    ) -> CompletionResult:
        payload: dict[str, Any] = {
            "model": model,
            "messages": [{"role": m.role, "content": m.content} for m in messages],
            "temperature": temperature,
        }
        if max_tokens is not None:
            payload["max_tokens"] = max_tokens
        if json_mode:
            payload["response_format"] = {"type": "json_object"}

        data = http_json_request(
            f"{self._base_url}/chat/completions",
            headers=self._headers(),
            payload=payload,
        )
        choice = (data.get("choices") or [{}])[0]
        text = (choice.get("message") or {}).get("content") or ""
        usage_raw = data.get("usage") or {}
        usage = Usage(
            prompt_tokens=int(usage_raw.get("prompt_tokens") or 0),
            completion_tokens=int(usage_raw.get("completion_tokens") or 0),
        )
        return CompletionResult(
            text=text.strip(),
            model=str(data.get("model") or model),
            provider=self.name,
            usage=usage,
        )

    def stream_complete(
        self,
        model: str,
        messages: list[Message],
        *,
        temperature: float = 0.7,
        max_tokens: int | None = None,
    ) -> Iterator[StreamChunk]:
        if not self.is_available():
            yield from default_stream_complete(
                self, model, messages, temperature=temperature, max_tokens=max_tokens
            )
            return

        payload: dict[str, Any] = {
            "model": model,
            "messages": [{"role": m.role, "content": m.content} for m in messages],
            "temperature": temperature,
            "stream": True,
        }
        if max_tokens is not None:
            payload["max_tokens"] = max_tokens

        usage: Usage | None = None
        for line in http_stream_request(
            f"{self._base_url}/chat/completions",
            headers=self._headers(),
            payload=payload,
        ):
            if not line.startswith("data:"):
                continue
            chunk = line[5:].strip()
            if chunk == "[DONE]":
                yield StreamChunk(text="", done=True, usage=usage)
                return
            try:
                parsed = json.loads(chunk)
            except json.JSONDecodeError:
                continue
            usage_raw = parsed.get("usage")
            if usage_raw:
                usage = Usage(
                    prompt_tokens=int(usage_raw.get("prompt_tokens") or 0),
                    completion_tokens=int(usage_raw.get("completion_tokens") or 0),
                )
            delta = ((parsed.get("choices") or [{}])[0].get("delta") or {}).get("content") or ""
            if delta:
                yield StreamChunk(text=delta, done=False)

        yield StreamChunk(text="", done=True, usage=usage)

    def embed(self, model: str, texts: list[str]) -> EmbeddingResult:
        data = http_json_request(
            f"{self._base_url}/embeddings",
            headers=self._headers(),
            payload={"model": model, "input": texts},
        )
        rows = sorted(data.get("data") or [], key=lambda x: x.get("index", 0))
        embeddings = [row.get("embedding") or [] for row in rows]
        return EmbeddingResult(
            embeddings=embeddings,
            model=str(data.get("model") or model),
            provider=self.name,
        )

    def rerank(
        self,
        model: str,
        query: str,
        documents: list[str],
        *,
        top_n: int | None = None,
    ) -> RerankResult:
        payload: dict[str, Any] = {
            "model": model,
            "query": query,
            "documents": documents,
            "return_documents": False,
        }
        if top_n is not None:
            payload["top_n"] = top_n
        data = http_json_request(
            f"{self._base_url}/rerank",
            headers=self._headers(),
            payload=payload,
        )
        rows = sorted(data.get("results") or [], key=lambda x: x.get("index", 0))
        scores = [float(row.get("relevance_score") or 0.0) for row in rows]
        indices = [int(row.get("index", i)) for i, row in enumerate(rows)]
        return RerankResult(
            scores=scores,
            indices=indices,
            model=str(data.get("model") or model),
            provider=self.name,
        )

    def ocr(
        self,
        model: str,
        image_b64: str,
        *,
        mime: str = "image/png",
        prompt: str = "请提取图片中的全部文字，保持段落结构。",
    ) -> OcrResult:
        data_url = f"data:{mime};base64,{image_b64}"
        payload: dict[str, Any] = {
            "model": model,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": data_url}},
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
            "temperature": 0.1,
            "max_tokens": 4096,
        }
        data = http_json_request(
            f"{self._base_url}/chat/completions",
            headers=self._headers(),
            payload=payload,
        )
        choice = (data.get("choices") or [{}])[0]
        text = (choice.get("message") or {}).get("content") or ""
        return OcrResult(
            text=text.strip(),
            model=str(data.get("model") or model),
            provider=self.name,
        )
