#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：django_base_ai
@File    ：base.py
@Desc    ：AI 客户端抽象基类与通用数据结构
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Union


@dataclass
class ModelItem:
    """自定义供应商下的单个模型。"""

    model_id: str  # 请求时使用的模型 ID，如 deepseek-chat
    name: str = ""  # 展示名称，如 DeepSeek Chat


@dataclass
class CustomProviderConfig:
    """
    自定义 API 配置，与前端「编辑供应商」表单对应。
    支持 API 格式（如 OpenAI）、前置 URL、模型列表等。
    """

    name: str
    api_key: str
    base_url: str  # API 前置 URL，如 https://api.deepseek.com/v1
    api_format: str = "OpenAI"  # 如 OpenAI
    models: list[ModelItem] = field(default_factory=list)
    default_model: str | None = None  # 默认使用的 model_id，不传则取 models[0].model_id


@dataclass
class ChatMessage:
    """单条对话消息。content 为 str 或多模态 list（如 [{"type":"text","text":"..."},{"type":"image_url","image_url":{"url":"data:..."}}]）。"""

    role: str  # system / user / assistant
    content: Union[str, list[Any]]  # 纯文本或 OpenAI 多模态 content 数组


@dataclass
class CompletionResult:
    """一次 completion 的返回结果。"""

    content: str
    model: str = ""
    usage: dict[str, int] = field(default_factory=dict)
    raw: Any = None
    finish_reason: str = ""


class ChatSession:
    """
    带会话上下文的对话封装，自动累积 user/assistant 消息，多轮交互保持上下文。
    """

    def __init__(
        self,
        client: "BaseAIClient",
        system_prompt: str | None = None,
        max_history_turns: int | None = None,
    ):
        """
        :param client: AI 客户端
        :param system_prompt: 可选系统提示，置于对话开头
        :param max_history_turns: 可选，保留最近 N 轮（每轮 user+assistant），超出则从最早开始丢弃，避免上下文过长
        """
        self._client = client
        self._messages: list[ChatMessage] = []
        if system_prompt:
            self._messages.append(ChatMessage(role="system", content=system_prompt))
        self._max_history_turns = max_history_turns

    @property
    def messages(self) -> list[ChatMessage]:
        """当前会话内的全部消息（含 system/user/assistant）。"""
        return self._messages.copy()

    def send(self, user_content: str, model: str | None = None, **kwargs: Any) -> CompletionResult:
        """
        发送用户内容，调用模型并自动将会话上下文（含本次 user 与 assistant 回复）写入会话。
        :param user_content: 用户本轮输入
        :param model: 可选模型名
        :param kwargs: 透传给 client.chat 的参数（如 temperature、max_tokens）
        :return: 本轮的 CompletionResult
        """
        self._messages.append(ChatMessage(role="user", content=user_content))
        if self._max_history_turns is not None:
            self._trim_to_max_turns()
        result = self._client.chat(self._messages, model=model, **kwargs)
        self._messages.append(ChatMessage(role="assistant", content=result.content))
        if self._max_history_turns is not None:
            self._trim_to_max_turns()
        return result

    def _trim_to_max_turns(self) -> None:
        """保留 system（若有）+ 最近 max_history_turns 轮（每轮 user+assistant）。"""
        if self._max_history_turns is None or self._max_history_turns <= 0:
            return
        system_msgs = [m for m in self._messages if m.role == "system"]
        rest = [m for m in self._messages if m.role != "system"]
        keep = self._max_history_turns * 2
        if len(rest) > keep:
            rest = rest[-keep:]
        self._messages = system_msgs + rest

    def clear(self) -> None:
        """清空会话上下文（若构造时带 system_prompt，会保留一条 system 消息）。"""
        system_msgs = [m for m in self._messages if m.role == "system"]
        self._messages = system_msgs.copy()


class BaseAIClient(ABC):
    """AI 调用客户端抽象基类，统一 chat/completion 接口。"""

    @abstractmethod
    def chat(
        self,
        messages: list[ChatMessage],
        model: str | None = None,
        **kwargs: Any,
    ) -> CompletionResult:
        """
        发送对话请求并返回补全结果。
        :param messages: 消息列表，通常含 system/user/assistant
        :param model: 模型名，由子类默认值或实现决定
        :param kwargs: 厂商扩展参数（temperature、max_tokens 等）
        :return: CompletionResult
        """
        pass

    def test_connection(self) -> tuple[bool, str]:
        """
        测试当前配置是否可用（如校验 API Key / 前置 URL）。
        :return: (是否成功, 消息)
        """
        try:
            self.chat([ChatMessage(role="user", content="hi")], max_tokens=2)
            return True, ""
        except Exception as e:
            return False, str(e)
