"""使用OpenAI兼容API进行AI推理的模型客户端。"""

import json
import re
import sys
import time
from dataclasses import dataclass, field
from typing import Any, Generator, Optional

from openai import OpenAI

from ubox_py_sdk.logger import get_logger

logger = get_logger("ubox_py_sdk.phone_agent.model")

# 类型定义
MessageDict = dict[str, Any]


@dataclass
class ModelConfig:
    """AI模型的配置类。"""

    base_url: str = "http://localhost:8000/v1"
    """模型服务的API基础URL。例如：http://localhost:8000/v1 或 https://api.openai.com/v1"""

    api_key: str = "EMPTY"
    """API密钥。用于认证模型服务请求。如果服务不需要认证，可以设置为"EMPTY"。"""

    model_name: str = "autoglm-phone-9b"
    """模型名称。指定要使用的AI模型，例如：autoglm-phone-9b、gpt-4-vision-preview等。"""

    max_tokens: int = 3000
    """最大生成token数。限制模型响应的最大长度，防止响应过长。"""

    temperature: float = 0.0
    """温度参数。控制模型输出的随机性，范围0-2。0表示最确定性的输出，值越大越随机。"""

    top_p: float = 0.85
    """Top-p采样参数。核采样（nucleus sampling）的概率阈值，范围0-1。只考虑累积概率达到此值的token。"""

    frequency_penalty: float = 0.2
    """频率惩罚。减少重复内容的概率，范围-2.0到2.0。正值会降低重复token的概率。"""

    extra_body: dict[str, Any] = field(default_factory=dict)
    """额外请求体参数。用于传递模型特定的额外配置参数，会合并到API请求的extra_body中。"""


@dataclass
class ModelResponse:
    """AI模型的响应。"""

    thinking: str
    action: str
    raw_content: str
    # 性能指标
    time_to_first_token: Optional[float] = None  # 首Token延迟（秒）
    time_to_thinking_end: Optional[float] = None  # 思考完成延迟（秒）
    total_time: Optional[float] = None  # 总推理时间（秒）
    # Token统计
    prompt_tokens: Optional[int] = None  # 输入token数
    completion_tokens: Optional[int] = None  # 输出token数
    total_tokens: Optional[int] = None  # 总token数


class ModelClient:
    """
    用于与OpenAI兼容的视觉语言模型交互的客户端。

    Args:
        config: 模型配置。
    """

    def __init__(self, config: Optional[ModelConfig] = None):
        self.config = config or ModelConfig()
        self.client = OpenAI(base_url=self.config.base_url, api_key=self.config.api_key)
        self._last_response: Optional[ModelResponse] = None

    def request(self, messages: list[MessageDict]) -> ModelResponse:
        """
        向模型发送请求。

        Args:
            messages: OpenAI格式的消息字典列表。

        Returns:
            包含思考和动作的ModelResponse。

        Raises:
            ValueError: 如果响应无法解析。
        """
        # 开始计时
        start_time = time.time()
        time_to_first_token = None
        time_to_thinking_end = None

        stream = self.client.chat.completions.create(
            messages=messages,
            model=self.config.model_name,
            max_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
            top_p=self.config.top_p,
            frequency_penalty=self.config.frequency_penalty,
            extra_body=self.config.extra_body,
            stream=True,
        )

        raw_content = ""
        buffer = ""  # 缓冲区，用于保存可能是标记一部分的内容
        action_markers = ["finish(message=", "do(action="]
        in_action_phase = False  # 跟踪是否已进入动作阶段
        first_token_received = False

        chunk_count = 0
        usage_info = None  # 用于保存token使用信息
        for chunk in stream:
            chunk_count += 1
            # 检查是否有usage信息（通常在最后一个chunk中）
            if hasattr(chunk, 'usage') and chunk.usage is not None:
                usage_info = chunk.usage

            if len(chunk.choices) == 0:
                continue
            if chunk.choices[0].delta.content is not None:
                content = chunk.choices[0].delta.content
                raw_content += content  # 所有内容都会被收集到 raw_content，包括思考过程

                # 记录首Token时间
                if not first_token_received:
                    time_to_first_token = time.time() - start_time
                    first_token_received = True

                if in_action_phase:
                    # 已在动作阶段，只累积内容不打印
                    continue

                buffer += content

                # 检查缓冲区中是否完整包含任何标记
                marker_found = False
                for marker in action_markers:
                    if marker in buffer:
                        # 找到标记，打印标记之前的所有内容（流式输出，不换行）
                        thinking_part = buffer.split(marker, 1)[0]
                        print(thinking_part, end="", flush=True, file=sys.stdout)
                        print()  # 思考完成后打印换行
                        in_action_phase = True
                        marker_found = True

                        # 记录思考完成时间
                        if time_to_thinking_end is None:
                            time_to_thinking_end = time.time() - start_time

                        break

                if marker_found:
                    continue  # 继续收集剩余内容

                # 检查缓冲区是否以任何标记的前缀结尾
                # 如果是，暂时不打印（等待更多内容）
                is_potential_marker = False
                for marker in action_markers:
                    for i in range(1, len(marker)):
                        if buffer.endswith(marker[:i]):
                            is_potential_marker = True
                            break
                    if is_potential_marker:
                        break

                if not is_potential_marker:
                    # 可以安全地打印缓冲区（流式输出，不换行）
                    print(buffer, end="", flush=True, file=sys.stdout)
                    buffer = ""

        # 计算总时间
        total_time = time.time() - start_time

        # 打印原始响应内容（包含所有内容，包括思考过程）
        logger.debug("🔍 [调试] 模型原始响应内容（完整）:")
        logger.debug(raw_content)
        logger.debug("=" * 80 + "\n")

        # 从响应中解析思考和动作
        thinking, action = self._parse_response(raw_content)

        # 提取token使用信息
        prompt_tokens = None
        completion_tokens = None
        total_tokens = None
        if usage_info:
            prompt_tokens = getattr(usage_info, 'prompt_tokens', None)
            completion_tokens = getattr(usage_info, 'completion_tokens', None)
            total_tokens = getattr(usage_info, 'total_tokens', None)

        # 打印性能指标
        logger.debug("")
        logger.debug("=" * 50)
        logger.debug("⏱️  性能指标:")
        logger.debug("-" * 50)
        if time_to_first_token is not None:
            logger.debug(
                f"首 Token 延迟 (TTFT): {time_to_first_token:.3f}s"
            )
        if time_to_thinking_end is not None:
            logger.debug(
                f"思考完成延迟:        {time_to_thinking_end:.3f}s"
            )
        logger.debug(
            f"总推理时间:          {total_time:.3f}s"
        )
        if total_tokens is not None:
            logger.debug("")
            logger.debug("📊 Token统计:")
            logger.debug("-" * 50)
            if prompt_tokens is not None:
                logger.debug(f"输入 Token:          {prompt_tokens}")
            if completion_tokens is not None:
                logger.debug(f"输出 Token:          {completion_tokens}")
            if total_tokens is not None:
                logger.debug(f"总 Token:            {total_tokens}")
        logger.debug("=" * 50)

        response = ModelResponse(
            thinking=thinking,
            action=action,
            raw_content=raw_content,
            time_to_first_token=time_to_first_token,
            time_to_thinking_end=time_to_thinking_end,
            total_time=total_time,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=total_tokens,
        )
        self._last_response = response
        return response

    def request_stream(self, messages: list[MessageDict]) -> Generator[str, None, None]:
        """
        向模型发送流式请求，实时返回思考过程。

        Args:
            messages: OpenAI格式的消息字典列表。

        Yields:
            str: 思考过程的文本块。
        """
        start_time = time.time()
        time_to_first_token = None
        time_to_thinking_end = None
        first_token_received = False

        stream = self.client.chat.completions.create(
            messages=messages,
            model=self.config.model_name,
            max_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
            top_p=self.config.top_p,
            frequency_penalty=self.config.frequency_penalty,
            extra_body=self.config.extra_body,
            stream=True,
        )

        raw_content = ""
        buffer = ""
        action_markers = ["finish(message=", "do(action="]
        in_action_phase = False
        usage_info = None  # 保存token使用信息

        for chunk in stream:
            if hasattr(chunk, 'usage') and chunk.usage is not None:
                usage_info = chunk.usage

            if len(chunk.choices) == 0:
                continue
            if chunk.choices[0].delta.content is not None:
                content = chunk.choices[0].delta.content
                raw_content += content

                if not first_token_received:
                    time_to_first_token = time.time() - start_time
                    first_token_received = True

                if in_action_phase:
                    continue

                buffer += content

                marker_found = False
                for marker in action_markers:
                    if marker in buffer:
                        thinking_part = buffer.split(marker, 1)[0]
                        if thinking_part:
                            yield thinking_part
                        in_action_phase = True
                        marker_found = True

                        if time_to_thinking_end is None:
                            time_to_thinking_end = time.time() - start_time
                        break

                if marker_found:
                    continue

                is_potential_marker = False
                for marker in action_markers:
                    for i in range(1, len(marker)):
                        if buffer.endswith(marker[:i]):
                            is_potential_marker = True
                            break
                    if is_potential_marker:
                        break

                if not is_potential_marker and buffer:
                    yield buffer
                    buffer = ""

        total_time = time.time() - start_time

        thinking, action = self._parse_response(raw_content)

        # 提取token使用信息
        prompt_tokens = None
        completion_tokens = None
        total_tokens = None
        if usage_info:
            prompt_tokens = getattr(usage_info, 'prompt_tokens', None)
            completion_tokens = getattr(usage_info, 'completion_tokens', None)
            total_tokens = getattr(usage_info, 'total_tokens', None)

        response = ModelResponse(
            thinking=thinking,
            action=action,
            raw_content=raw_content,
            time_to_first_token=time_to_first_token,
            time_to_thinking_end=time_to_thinking_end,
            total_time=total_time,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=total_tokens,
        )
        self._last_response = response

        logger.debug("")
        logger.debug("=" * 50)
        logger.debug("⏱️  性能指标:")
        logger.debug("-" * 50)
        if time_to_first_token is not None:
            logger.debug(f"首 Token 延迟 (TTFT): {time_to_first_token:.3f}s")
        if time_to_thinking_end is not None:
            logger.debug(f"思考完成延迟:        {time_to_thinking_end:.3f}s")
        logger.debug(f"总推理时间:          {total_time:.3f}s")
        if total_tokens is not None:
            logger.debug("")
            logger.debug("📊 Token统计:")
            logger.debug("-" * 50)
            if prompt_tokens is not None:
                logger.debug(f"输入 Token:          {prompt_tokens}")
            if completion_tokens is not None:
                logger.debug(f"输出 Token:          {completion_tokens}")
            if total_tokens is not None:
                logger.debug(f"总 Token:            {total_tokens}")
        logger.debug("=" * 50)

    def get_last_response(self) -> Optional[ModelResponse]:
        """获取最后一次请求的完整响应。"""
        return self._last_response

    def _parse_response(self, content: str) -> tuple[str, str]:
        """
        将模型响应解析为思考和动作部分。

        解析规则：
        1. 如果内容包含 'finish(message='，之前的所有内容为思考部分，
           从 'finish(message=' 开始的所有内容为动作部分。
        2. 如果规则1不适用但内容包含 'do(action='，
           之前的所有内容为思考部分，从 'do(action=' 开始的所有内容为动作部分。
        3. 回退：如果内容包含 '<answer>'，使用传统的XML标签解析。
        4. 否则，返回空的思考部分和完整内容作为动作。

        Args:
            content: 原始响应内容。

        Returns:
            (思考, 动作) 元组。
        """
        # 规则1：检查 finish(message=
        if "finish(message=" in content:
            parts = content.split("finish(message=", 1)
            thinking = parts[0].strip()
            action = "finish(message=" + parts[1]
            return thinking, action

        # 规则2：检查 do(action=
        if "do(action=" in content:
            parts = content.split("do(action=", 1)
            thinking = parts[0].strip()
            action = "do(action=" + parts[1]
            return thinking, action

        # 规则3：回退到传统的XML标签解析
        # 处理 <think> 和 <answer> 可能独立出现的情况
        thinking = ""
        action = ""

        # 提取 <think> 内容
        if "<think>" in content:
            # 使用正则表达式提取 <think>...</think> 的内容
            thinking_match = re.search(r'<think>(.*?)</think>', content, re.DOTALL)
            if thinking_match:
                thinking = thinking_match.group(1).strip()

        # 提取 <answer> 内容
        if "<answer>" in content:
            # 使用正则表达式提取 <answer>...</answer> 的内容
            answer_match = re.search(r'<answer>(.*?)</answer>', content, re.DOTALL)
            if answer_match:
                action = answer_match.group(1).strip()

        # 如果找到了 thinking 或 action，返回它们
        if thinking or action:
            return thinking, action

        # 规则4：未找到标记，返回内容作为动作
        return "", content


class MessageBuilder:
    """用于构建对话消息的辅助类。"""

    @staticmethod
    def create_system_message(content: str) -> dict[str, Any]:
        """创建系统消息。"""
        return {"role": "system", "content": content}

    @staticmethod
    def create_user_message(
            text: str, image_base64: Optional[str] = None
    ) -> dict[str, Any]:
        """
        创建用户消息（可选包含图片）。

        Args:
            text: 文本内容。
            image_base64: 可选的base64编码图片。

        Returns:
            消息字典。
        """
        content = []

        if image_base64:
            content.append(
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{image_base64}"},
                }
            )

        content.append({"type": "text", "text": text})

        return {"role": "user", "content": content}

    @staticmethod
    def create_assistant_message(content: str) -> MessageDict:
        """
        创建助手消息。

        Args:
            content: 助手消息内容。

        Returns:
            助手消息字典。
        """
        return {"role": "assistant", "content": content}

    @staticmethod
    def remove_images_from_message(message: MessageDict) -> MessageDict:
        """
        从消息中移除图片内容以节省上下文空间。

        Args:
            message: 消息字典。

        Returns:
            移除图片后的消息。
        """
        if isinstance(message.get("content"), list):
            message["content"] = [
                item for item in message["content"] if item.get("type") == "text"
            ]
        return message

    @staticmethod
    def build_screen_info(current_app: str, **extra_info) -> str:
        """
        为模型构建屏幕信息字符串。

        Args:
            current_app: 当前应用名称。
            **extra_info: 要包含的额外信息。

        Returns:
            包含屏幕信息的JSON字符串。
        """
        info = {"current_app": current_app, **extra_info}
        return json.dumps(info, ensure_ascii=False)
