from __future__ import annotations

from typing import Any, Iterator, Protocol, runtime_checkable

from niannian.llm.types import CompletionResult, EmbeddingResult, Message, OcrResult, RerankResult, StreamChunk


@runtime_checkable
class LLMProvider(Protocol):
    name: str

    def is_available(self) -> bool:
        """Provider 是否已配置可用（如 API Key 存在）。"""

    def complete(
        self,
        model: str,
        messages: list[Message],
        *,
        temperature: float = 0.7,
        max_tokens: int | None = None,
        json_mode: bool = False,
    ) -> CompletionResult:
        """非流式补全。"""

    def stream_complete(
        self,
        model: str,
        messages: list[Message],
        *,
        temperature: float = 0.7,
        max_tokens: int | None = None,
    ) -> Iterator[StreamChunk]:
        """流式补全；默认回退为非流式单次输出。"""

    def embed(self, model: str, texts: list[str]) -> EmbeddingResult:
        """向量嵌入。"""

    def rerank(
        self,
        model: str,
        query: str,
        documents: list[str],
        *,
        top_n: int | None = None,
    ) -> RerankResult:
        """向量重排。"""

    def ocr(
        self,
        model: str,
        image_b64: str,
        *,
        mime: str = "image/png",
        prompt: str = "请提取图片中的全部文字，保持段落结构。",
    ) -> OcrResult:
        """视觉 OCR。"""


def default_stream_complete(
    provider: LLMProvider,
    model: str,
    messages: list[Message],
    *,
    temperature: float = 0.7,
    max_tokens: int | None = None,
) -> Iterator[StreamChunk]:
    result = provider.complete(
        model,
        messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    yield StreamChunk(text=result.text, done=True, usage=result.usage)


def http_json_request(
    url: str,
    *,
    headers: dict[str, str],
    payload: dict[str, Any],
    timeout: float = 120.0,
) -> dict[str, Any]:
    import json
    import urllib.error
    import urllib.request

    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            body = resp.read().decode("utf-8")
            return json.loads(body) if body else {}
    except urllib.error.HTTPError as e:
        err_body = e.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {e.code}: {err_body[:500]}") from e


def http_stream_request(
    url: str,
    *,
    headers: dict[str, str],
    payload: dict[str, Any],
    timeout: float = 120.0,
) -> Iterator[str]:
    import json
    import urllib.error
    import urllib.request

    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            for raw_line in resp:
                line = raw_line.decode("utf-8", errors="replace").strip()
                if line:
                    yield line
    except urllib.error.HTTPError as e:
        err_body = e.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {e.code}: {err_body[:500]}") from e
