from __future__ import annotations

import asyncio
import typing
import json
import logging
import tempfile
import os

from tboxsdk.tbox import TboxClient
from tboxsdk.model.file import File, FileType

from .. import runner
from ...core import app
from ...utils import bounded_executor, image
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message

_MAX_TBOX_RESPONSE_CHARS = 1024 * 1024
_MAX_TBOX_MEDIA_BYTES = 10 * 1024 * 1024


class TboxAPIError(Exception):
    """TBox API 请求失败"""

    def __init__(self, message: str):
        self.message = message
        super().__init__(self.message)


def _append_bounded(current: str, addition: typing.Any) -> str:
    addition = str(addition or '')
    if len(current) + len(addition) > _MAX_TBOX_RESPONSE_CHARS:
        raise TboxAPIError('Tbox response exceeds the runtime limit')
    return current + addition


def _write_temp_media(file_bytes: bytes, suffix: str) -> str:
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp_file:
        tmp_file.write(file_bytes)
        return tmp_file.name


@runner.runner_class('tbox-app-api')
class TboxAPIRunner(runner.RequestRunner):
    "蚂蚁百宝箱API对话请求器"

    # 运行器内部使用的配置
    app_id: str  # 蚂蚁百宝箱平台中的应用ID
    api_key: str  # 在蚂蚁百宝箱平台中申请的令牌

    def __init__(self, ap: app.Application, pipeline_config: dict):
        """初始化"""
        self.ap = ap
        self.pipeline_config = pipeline_config

        # 初始化Tbox 参数配置
        self.app_id = self.pipeline_config['ai']['tbox-app-api']['app-id']
        self.api_key = self.pipeline_config['ai']['tbox-app-api']['api-key']

        # 初始化Tbox client
        logging.getLogger('tbox.client').setLevel(logging.WARNING)
        self.tbox_client = TboxClient(authorization=self.api_key)

    async def _preprocess_user_message(self, query: pipeline_query.Query) -> tuple[str, list[str]]:
        """预处理用户消息，提取纯文本，并将图片上传到 Tbox 服务

        Returns:
            tuple[str, list[str]]: 纯文本和图片的 Tbox 文件ID
        """
        plain_text = ''
        image_ids = []

        if isinstance(query.user_message.content, list):
            for ce in query.user_message.content:
                if ce.type == 'text':
                    plain_text += ce.text
                elif ce.type == 'image_base64':
                    image_b64, image_format = await image.extract_b64_and_format(ce.image_base64)
                    file_bytes = await image.decode_base64_limited(
                        image_b64,
                        max_bytes=_MAX_TBOX_MEDIA_BYTES,
                    )
                    tmp_file_path: str | None = None
                    try:
                        tmp_file_path = await asyncio.to_thread(
                            _write_temp_media,
                            file_bytes,
                            f'.{image_format}',
                        )
                        file_upload_resp = await asyncio.to_thread(
                            self.tbox_client.upload_file,
                            tmp_file_path,
                        )
                        image_id = file_upload_resp.get('data', '')
                        image_ids.append(image_id)
                    finally:
                        if tmp_file_path and os.path.exists(tmp_file_path):
                            await bounded_executor.run_blocking_cleanup(
                                os.unlink,
                                tmp_file_path,
                            )
        elif isinstance(query.user_message.content, str):
            plain_text = query.user_message.content

        return plain_text, image_ids

    async def _agent_messages(
        self, query: pipeline_query.Query
    ) -> typing.AsyncGenerator[provider_message.Message, None]:
        """TBox 智能体对话请求"""

        plain_text, image_ids = await self._preprocess_user_message(query)
        remove_think = self.pipeline_config['output'].get('misc', {}).get('remove-think')

        try:
            is_stream = await query.adapter.is_stream_output_supported()
        except AttributeError:
            is_stream = False

        # 获取Tbox的conversation_id
        conversation_id = query.session.using_conversation.uuid or None

        files = None
        if image_ids:
            files = [File(file_id=image_id, type=FileType.IMAGE) for image_id in image_ids]

        # 发送对话请求
        response = await asyncio.to_thread(
            self.tbox_client.chat,
            app_id=self.app_id,
            user_id=query.bot_uuid,
            query=plain_text,
            stream=is_stream,
            conversation_id=conversation_id,
            files=files,
        )

        if is_stream:
            # 解析Tbox流式输出内容，并发送给上游
            async for chunk in self._process_stream_message(
                response,
                query,
                remove_think,
            ):
                yield chunk
        else:
            message = self._process_non_stream_message(response, query, remove_think)
            yield provider_message.Message(
                role='assistant',
                content=message,
            )

    def _process_non_stream_message(self, response: typing.Dict, query: pipeline_query.Query, remove_think: bool):
        if response.get('errorCode') != '0':
            raise TboxAPIError(f'Tbox API 请求失败: {response.get("errorMsg", "")}')
        payload = response.get('data', {})
        conversation_id = payload.get('conversationId', '')
        query.session.using_conversation.uuid = conversation_id
        thinking_content = payload.get('reasoningContent', [])
        result = ''
        if thinking_content and not remove_think:
            result = _append_bounded(
                result,
                f'<think>\n{thinking_content[0].get("text", "")}\n</think>\n',
            )
        content = payload.get('result', [])
        if content:
            result = _append_bounded(result, content[0].get('chunk', ''))
        return result

    async def _process_stream_message(
        self, response: typing.Generator[dict], query: pipeline_query.Query, remove_think: bool
    ):
        idx_msg = 0
        pending_content = ''
        conversation_id = None
        think_start = False
        think_end = False
        async for chunk in runner.iterate_sync(response):
            if chunk.get('type', '') == 'chunk':
                """
                Tbox返回的消息内容chunk结构
                {'lane': 'default', 'payload': {'conversationId': '20250918tBI947065406', 'messageId': '20250918TB1f53230954', 'text': '️'}, 'type': 'chunk'}
                """
                # 如果包含思考过程，拼接</think>
                if think_start and not think_end:
                    pending_content = _append_bounded(
                        pending_content,
                        '\n</think>\n',
                    )
                    think_end = True

                payload = chunk.get('payload', {})
                if not conversation_id:
                    conversation_id = payload.get('conversationId')
                    query.session.using_conversation.uuid = conversation_id
                if payload.get('text'):
                    idx_msg += 1
                    pending_content = _append_bounded(
                        pending_content,
                        payload.get('text'),
                    )
            elif chunk.get('type', '') == 'thinking' and not remove_think:
                """
                Tbox返回的思考过程chunk结构
                {'payload': '{"ext_data":{"text":"日期"},"event":"flow.node.llm.thinking","entity":{"node_type":"text-completion","execute_id":"6","group_id":0,"parent_execute_id":"6","node_name":"模型推理","node_id":"TC_5u6gl0"}}', 'type': 'thinking'}
                """
                payload = json.loads(chunk.get('payload', '{}'))
                if payload.get('ext_data', {}).get('text'):
                    idx_msg += 1
                    content = payload.get('ext_data', {}).get('text')
                    if not think_start:
                        think_start = True
                        pending_content = _append_bounded(
                            pending_content,
                            f'<think>\n{content}',
                        )
                    else:
                        pending_content = _append_bounded(
                            pending_content,
                            content,
                        )
            elif chunk.get('type', '') == 'error':
                raise TboxAPIError(
                    f'Tbox API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} '
                )

            if idx_msg % 8 == 0:
                yield provider_message.MessageChunk(
                    role='assistant',
                    content=pending_content,
                    is_final=False,
                )

        # Tbox不返回END事件，默认发一个最终消息
        yield provider_message.MessageChunk(
            role='assistant',
            content=pending_content,
            is_final=True,
        )

    async def run(self, query: pipeline_query.Query) -> typing.AsyncGenerator[provider_message.Message, None]:
        """运行"""
        msg_seq = 0
        async for msg in self._agent_messages(query):
            if isinstance(msg, provider_message.MessageChunk):
                msg_seq += 1
                msg.msg_sequence = msg_seq
            yield msg
