"""智谱 GLM 适配器 -- OpenAI 兼容格式"""

from __future__ import annotations

import json
from typing import Any, AsyncGenerator

import httpx

from app.providers.base import BaseProvider, Capability, QuotaInfo


class ZhipuProvider(BaseProvider):
    name = "zhipu"
    display_name = "智谱 GLM"
    capabilities = [Capability.CHAT, Capability.IMAGE]

    async def chat(
        self,
        messages: list[dict],
        model: str,
        plan: dict,
        stream: bool = True,
        **kwargs,
    ) -> AsyncGenerator[str, None]:
        url = f"{self._base_url(plan)}/chat/completions"
        body: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": stream,
            **kwargs,
        }
        headers = self._build_headers(plan)

        async with httpx.AsyncClient(timeout=120) as client:
            if stream:
                async with client.stream("POST", url, json=body, headers=headers) as resp:
                    resp.raise_for_status()
                    async for line in resp.aiter_lines():
                        if line.startswith("data: "):
                            yield line + "\n\n"
            else:
                resp = await client.post(url, json=body, headers=headers)
                resp.raise_for_status()
                yield json.dumps(resp.json())

    async def generate_image(self, prompt: str, plan: dict, **kwargs) -> dict[str, Any]:
        url = f"{self._base_url(plan)}/images/generations"
        body = {"model": kwargs.get("model", "cogview-3"), "prompt": prompt, **kwargs}
        headers = self._build_headers(plan)
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(url, json=body, headers=headers)
            resp.raise_for_status()
            return resp.json()

    async def query_quota(self, plan: dict) -> QuotaInfo | None:
        """智谱余额查询"""
        try:
            url = f"{self._base_url(plan)}/../dashboard/billing/usage"
            headers = self._build_headers(plan)
            async with httpx.AsyncClient(timeout=30) as client:
                resp = await client.get(url, headers=headers)
                resp.raise_for_status()
                data = resp.json()
                return QuotaInfo(raw=data)
        except Exception:
            return None