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, StreamChunk, Usage


class AnthropicProvider:
    """Anthropic Messages API。"""

    name = "anthropic"

    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.anthropic.com"
        self._base_url = str(base).rstrip("/")
        self._api_version = str(config.get("api_version") or "2023-06-01")

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

    def _headers(self, *, stream: bool = False) -> dict[str, str]:
        headers = {
            "x-api-key": self._api_key,
            "anthropic-version": self._api_version,
            "Content-Type": "application/json",
        }
        if stream:
            headers["Accept"] = "text/event-stream"
        return headers

    @staticmethod
    def _split_system(messages: list[Message]) -> tuple[str | None, list[dict[str, str]]]:
        system_parts: list[str] = []
        out: list[dict[str, str]] = []
        for m in messages:
            if m.role == "system":
                system_parts.append(m.content)
            else:
                out.append({"role": m.role, "content": m.content})
        system = "\n\n".join(system_parts) if system_parts else None
        return system, out

    def complete(
        self,
        model: str,
        messages: list[Message],
        *,
        temperature: float = 0.7,
        max_tokens: int | None = None,
        json_mode: bool = False,
    ) -> CompletionResult:
        system, api_messages = self._split_system(messages)
        payload: dict[str, Any] = {
            "model": model,
            "messages": api_messages,
            "max_tokens": max_tokens or 4096,
            "temperature": temperature,
        }
        if system:
            payload["system"] = system
        if json_mode:
            payload["system"] = (payload.get("system") or "") + "\n\nRespond with valid JSON only."

        data = http_json_request(
            f"{self._base_url}/v1/messages",
            headers=self._headers(),
            payload=payload,
        )
        text_parts: list[str] = []
        for block in data.get("content") or []:
            if block.get("type") == "text":
                text_parts.append(block.get("text") or "")
        usage_raw = data.get("usage") or {}
        usage = Usage(
            prompt_tokens=int(usage_raw.get("input_tokens") or 0),
            completion_tokens=int(usage_raw.get("output_tokens") or 0),
        )
        return CompletionResult(
            text="".join(text_parts).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

        system, api_messages = self._split_system(messages)
        payload: dict[str, Any] = {
            "model": model,
            "messages": api_messages,
            "max_tokens": max_tokens or 4096,
            "temperature": temperature,
            "stream": True,
        }
        if system:
            payload["system"] = system

        usage: Usage | None = None
        for line in http_stream_request(
            f"{self._base_url}/v1/messages",
            headers=self._headers(stream=True),
            payload=payload,
        ):
            if not line.startswith("data:"):
                continue
            chunk = line[5:].strip()
            if not chunk:
                continue
            try:
                parsed = json.loads(chunk)
            except json.JSONDecodeError:
                continue
            event_type = parsed.get("type")
            if event_type == "content_block_delta":
                delta = (parsed.get("delta") or {}).get("text") or ""
                if delta:
                    yield StreamChunk(text=delta, done=False)
            elif event_type == "message_delta":
                usage_raw = (parsed.get("usage") or {})
                if usage_raw:
                    usage = Usage(
                        prompt_tokens=int(usage_raw.get("input_tokens") or 0),
                        completion_tokens=int(usage_raw.get("output_tokens") or 0),
                    )
            elif event_type == "message_stop":
                yield StreamChunk(text="", done=True, usage=usage)
                return

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

    def embed(self, model: str, texts: list[str]) -> EmbeddingResult:
        raise NotImplementedError("anthropic provider 不支持 embed；请使用 openai_compat embedding slot")

    def rerank(self, model: str, query: str, documents: list[str], *, top_n: int | None = None):
        raise NotImplementedError("anthropic provider 不支持 rerank；请使用 retrieval_rerank slot")

    def ocr(self, model: str, image_b64: str, *, mime: str = "image/png", prompt: str = ""):
        raise NotImplementedError("anthropic provider 不支持 ocr；请使用 api.ocr slot")
