#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：django_base_ai
@File    ：openai_compat.py
@Desc    ：OpenAI 兼容 API 客户端，适用于 DeepSeek、OpenAI 等
"""

import json
import logging
from typing import Any, Iterator

import requests

from django_base_ai.utils.ai.base import BaseAIClient, ChatMessage, CompletionResult

logger = logging.getLogger(__name__)

# 常见兼容端点（可按需扩展）
DEFAULT_BASE_URLS = {
    "openai": "https://api.openai.com/v1",
    "deepseek": "https://api.deepseek.com/v1",
    "openai_custom": "https://api.openai.com/v1",  # 自建或代理
}


class OpenAICompatClient(BaseAIClient):
    """
    OpenAI 兼容的 Chat Completions API 客户端。
    适用于：OpenAI、DeepSeek 官方等提供 OpenAI 格式的厂商。
    """

    def __init__(
        self,
        api_key: str,
        base_url: str | None = None,
        default_model: str = "gpt-3.5-turbo",
        provider: str = "openai_custom",
        timeout: int = 60,
    ):
        """
        :param api_key: API Key
        :param base_url: 兼容接口根地址，如 https://api.deepseek.com/v1
        :param default_model: 默认模型名
        :param provider: 预设名，用于从 DEFAULT_BASE_URLS 取 base_url（当 base_url 未传时）
        :param timeout: 请求超时秒数
        """
        self.api_key = api_key
        self.base_url = (base_url or DEFAULT_BASE_URLS.get(provider, "")).rstrip("/")
        self.default_model = default_model
        self.timeout = timeout
        self._session = requests.Session()
        self._session.headers.update(
            {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            }
        )

    def chat(
        self,
        messages: list[ChatMessage],
        model: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 2048,
        **kwargs: Any,
    ) -> CompletionResult:
        url = f"{self.base_url}/chat/completions"
        def _msg_content(m):
            if isinstance(m.content, list):
                return m.content
            return m.content

        body = {
            "model": model or self.default_model,
            "messages": [{"role": m.role, "content": _msg_content(m)} for m in messages],
            "temperature": temperature,
            "max_tokens": max_tokens,
            **{k: v for k, v in kwargs.items() if v is not None},
        }
        try:
            resp = self._session.post(url, json=body, timeout=self.timeout)
            resp.raise_for_status()
            data = resp.json()
        except requests.RequestException as e:
            logger.exception("OpenAI compat request failed: %s", e)
            raise
        choice = (data.get("choices") or [None])[0]
        if not choice:
            raise ValueError("API 返回无 choices")
        delta = choice.get("message", {})
        usage = data.get("usage") or {}
        return CompletionResult(
            content=delta.get("content", ""),
            model=data.get("model", model or self.default_model),
            usage={
                "prompt_tokens": usage.get("prompt_tokens", 0),
                "completion_tokens": usage.get("completion_tokens", 0),
                "total_tokens": usage.get("total_tokens", 0),
            },
            raw=data,
            finish_reason=choice.get("finish_reason", ""),
        )

    def chat_stream(
        self,
        messages: list[ChatMessage],
        model: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 2048,
        **kwargs: Any,
    ) -> Iterator[dict[str, Any]]:
        """
        流式对话，逐 chunk 产出，便于前端 EventSource 打字机效果。
        每次 yield: {"content": "delta"}；结束时 yield {"content": "", "done": True, "usage": {...}, "model": "..."}。
        """
        url = f"{self.base_url}/chat/completions"
        def _msg_content(m):
            if isinstance(m.content, list):
                return m.content
            return m.content

        body = {
            "model": model or self.default_model,
            "messages": [{"role": m.role, "content": _msg_content(m)} for m in messages],
            "stream": True,
            "temperature": temperature,
            "max_tokens": max_tokens,
            **{k: v for k, v in kwargs.items() if v is not None},
        }
        try:
            resp = self._session.post(
                url, json=body, timeout=self.timeout, stream=True
            )
            resp.raise_for_status()
        except requests.RequestException as e:
            logger.exception("OpenAI compat stream request failed: %s", e)
            raise
        model_used = model or self.default_model
        usage: dict[str, int] = {}
        try:
            for line in resp.iter_lines(decode_unicode=True):
                if not line or not line.strip():
                    continue
                if line.strip().startswith("data:"):
                    payload = line.strip()[5:].strip()
                    if payload == "[DONE]":
                        break
                    try:
                        data = json.loads(payload)
                    except json.JSONDecodeError:
                        continue
                    choices = data.get("choices") or []
                    if not choices:
                        usage = (data.get("usage") or {}).copy()
                        if data.get("model"):
                            model_used = data["model"]
                        continue
                    choice = choices[0]
                    delta = choice.get("delta") or {}
                    content = delta.get("content") or ""
                    if content:
                        yield {"content": content}
                    finish = choice.get("finish_reason")
                    if finish:
                        usage = (data.get("usage") or {}).copy()
                        if data.get("model"):
                            model_used = data["model"]
                        break
        finally:
            resp.close()
        yield {
            "content": "",
            "done": True,
            "usage": {
                "prompt_tokens": usage.get("prompt_tokens", 0),
                "completion_tokens": usage.get("completion_tokens", 0),
                "total_tokens": usage.get("total_tokens", 0),
            },
            "model": model_used,
        }
