from __future__ import annotations

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

import httpx

from ..config import CloudProfile, ModelConfig, RuntimeConfig
from ..types import BackendError, ChatChunk, ChatParams
from .base import BaseBackend


class RemoteBackend(BaseBackend):
    """Proxy any OpenAI-compatible cloud endpoint without exposing its key to clients."""

    def __init__(self, model: ModelConfig, profile: CloudProfile, runtime: RuntimeConfig) -> None:
        super().__init__(model, runtime)
        self.profile = profile
        self.base_url = profile.base_url.rstrip("/")
        if not self.base_url.endswith("/v1"):
            self.base_url += "/v1"

    async def load(self) -> None:
        key = self.profile.resolved_api_key()
        if not key:
            raise BackendError(
                f"云端配置 {self.profile.id} 没有 API key。请设置环境变量 {self.profile.api_key_env or '<api_key_env>'}。"
            )
        self.loaded = True

    async def unload(self) -> None:
        self.loaded = False

    def _headers(self) -> dict[str, str]:
        key = self.profile.resolved_api_key()
        if not key:
            raise BackendError(f"云端配置 {self.profile.id} 的 API key 不可用")
        headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}
        headers.update(self.profile.headers)
        return headers

    async def chat(
        self,
        messages: list[dict[str, Any]],
        params: ChatParams,
        stream: bool = False,
    ) -> AsyncIterator[ChatChunk]:
        if not self.loaded:
            await self.load()
        payload: dict[str, Any] = {
            "model": self.model.remote_model or self.profile.default_model or self.model.id,
            "messages": messages,
            "stream": stream,
            "max_tokens": params.max_tokens,
            "temperature": params.temperature,
            "top_p": params.top_p,
        }
        if params.stop:
            payload["stop"] = params.stop
        if params.seed is not None:
            payload["seed"] = params.seed
        if params.response_format is not None:
            payload["response_format"] = (
                {"type": params.response_format} if isinstance(params.response_format, str) else params.response_format
            )
        if params.tools:
            payload["tools"] = params.tools
        if params.tool_choice is not None:
            payload["tool_choice"] = params.tool_choice
        payload.update(params.extra)
        timeout = httpx.Timeout(self.profile.timeout, connect=15.0)
        async with httpx.AsyncClient(timeout=timeout, headers=self._headers()) as client:
            if stream:
                async with client.stream("POST", f"{self.base_url}/chat/completions", json=payload) as response:
                    if response.status_code >= 400:
                        detail = (await response.aread()).decode("utf-8", errors="replace")
                        raise BackendError(f"云端请求失败 {response.status_code}: {detail}")
                    async for line in response.aiter_lines():
                        if not line.startswith("data:"):
                            continue
                        raw = line[5:].strip()
                        if raw == "[DONE]":
                            yield ChatChunk(done=True, finish_reason="stop")
                            return
                        try:
                            data = json.loads(raw)
                        except json.JSONDecodeError:
                            continue
                        choice = (data.get("choices") or [{}])[0]
                        delta = choice.get("delta") or {}
                        text = delta.get("content") or ""
                        finish = choice.get("finish_reason")
                        if text or finish or len(delta) > 1:
                            yield ChatChunk(text=text, finish_reason=finish, metadata=data, delta=delta)
                    yield ChatChunk(done=True, finish_reason="stop")
                return

            response = await client.post(f"{self.base_url}/chat/completions", json=payload)
            if response.status_code >= 400:
                raise BackendError(f"云端请求失败 {response.status_code}: {response.text}")
            data = response.json()
            choice = (data.get("choices") or [{}])[0]
            message = choice.get("message") or {}
            yield ChatChunk(
                text=message.get("content") or "",
                done=True,
                finish_reason=choice.get("finish_reason") or "stop",
                metadata=data,
                delta=message,
            )

    def info(self) -> dict[str, Any]:
        return {
            "id": self.model.id,
            "kind": "cloud",
            "profile": self.profile.id,
            "remote_model": self.model.remote_model or self.profile.default_model or self.model.id,
            "loaded": self.loaded,
            "base_url": self.base_url,
        }