"""AI Provider 抽象层，支持 OpenAI 兼容 API 和 Anthropic 原生 API。"""

import json
from typing import AsyncGenerator

import httpx

from ..models import AIConfig


class StreamEvent:
    """流式事件：文本chunk或usage信息"""
    def __init__(self, text: str = "", usage: dict | None = None):
        self.text = text
        self.usage = usage


async def stream_chat(
    config: AIConfig,
    messages: list[dict],
    system_prompt: str = "",
) -> AsyncGenerator[StreamEvent, None]:
    """根据 provider 类型调用对应的流式 API，返回 StreamEvent。"""
    if config.provider == "anthropic":
        async for event in _stream_anthropic(config, messages, system_prompt):
            yield event
    else:
        async for event in _stream_openai(config, messages, system_prompt):
            yield event


async def _stream_openai(
    config: AIConfig,
    messages: list[dict],
    system_prompt: str,
) -> AsyncGenerator[StreamEvent, None]:
    """调用 OpenAI 兼容 API（stream=true）。"""
    url = f"{config.base_url.rstrip('/')}/chat/completions"

    payload_messages = []
    if system_prompt:
        payload_messages.append({"role": "system", "content": system_prompt})
    payload_messages.extend(messages)

    body = {
        "model": config.model,
        "messages": payload_messages,
        "stream": True,
        "stream_options": {"include_usage": True},
    }

    async with httpx.AsyncClient(timeout=120) as client:
        async with client.stream(
            "POST",
            url,
            headers={
                "Authorization": f"Bearer {config.api_key}",
                "Content-Type": "application/json",
            },
            json=body,
        ) as resp:
            if resp.status_code >= 400:
                body_text = await resp.aread()
                try:
                    err = json.loads(body_text)
                    msg = err.get("error", {}).get("message", "") or str(err)
                except Exception:
                    msg = body_text.decode(errors="replace")[:200]
                raise RuntimeError(f"API错误 ({resp.status_code}): {msg}")

            ct = resp.headers.get("content-type", "")
            if "text/html" in ct:
                raise RuntimeError(f"API地址返回了HTML页面，请检查API地址是否正确（当前: {url}）")

            async for line in resp.aiter_lines():
                if not line.startswith("data: "):
                    continue
                data = line[6:]
                if data.strip() == "[DONE]":
                    break
                try:
                    chunk = json.loads(data)
                    # 提取usage（OpenAI在最后一个chunk返回）
                    usage = chunk.get("usage")
                    if usage:
                        yield StreamEvent(usage={
                            "prompt_tokens": usage.get("prompt_tokens", 0),
                            "completion_tokens": usage.get("completion_tokens", 0),
                            "total_tokens": usage.get("total_tokens", 0),
                        })
                    choices = chunk.get("choices", [])
                    if choices:
                        delta = choices[0].get("delta", {})
                        content = delta.get("content")
                        if content:
                            yield StreamEvent(text=content)
                except (json.JSONDecodeError, KeyError, IndexError):
                    continue


async def _stream_anthropic(
    config: AIConfig,
    messages: list[dict],
    system_prompt: str,
) -> AsyncGenerator[StreamEvent, None]:
    """调用 Anthropic Messages API（stream=true）。"""
    url = f"{config.base_url.rstrip('/')}/messages"

    body: dict = {
        "model": config.model,
        "max_tokens": 4096,
        "messages": messages,
        "stream": True,
    }
    if system_prompt:
        body["system"] = system_prompt

    async with httpx.AsyncClient(timeout=120) as client:
        async with client.stream(
            "POST",
            url,
            headers={
                "x-api-key": config.api_key,
                "anthropic-version": "2023-06-01",
                "content-type": "application/json",
            },
            json=body,
        ) as resp:
            if resp.status_code >= 400:
                body_text = await resp.aread()
                try:
                    err = json.loads(body_text)
                    msg = err.get("error", {}).get("message", "") or str(err)
                except Exception:
                    msg = body_text.decode(errors="replace")[:200]
                raise RuntimeError(f"API错误 ({resp.status_code}): {msg}")

            ct = resp.headers.get("content-type", "")
            if "text/html" in ct:
                raise RuntimeError(f"API地址返回了HTML页面，请检查API地址是否正确（当前: {url}）")

            async for line in resp.aiter_lines():
                if not line.startswith("data: "):
                    continue
                try:
                    event = json.loads(line[6:])
                    if event.get("type") == "content_block_delta":
                        text = event.get("delta", {}).get("text", "")
                        if text:
                            yield StreamEvent(text=text)
                    elif event.get("type") == "message_delta":
                        usage = event.get("usage", {})
                        if usage:
                            yield StreamEvent(usage={
                                "prompt_tokens": 0,
                                "completion_tokens": usage.get("output_tokens", 0),
                                "total_tokens": 0,
                            })
                    elif event.get("type") == "message_start":
                        msg = event.get("message", {})
                        usage = msg.get("usage", {})
                        if usage:
                            yield StreamEvent(usage={
                                "prompt_tokens": usage.get("input_tokens", 0),
                                "completion_tokens": 0,
                                "total_tokens": 0,
                            })
                    elif event.get("type") == "message_stop":
                        break
                except json.JSONDecodeError:
                    continue