"""
企业微信智能机器人服务（WebSocket 长连接模式）

使用官方 SDK (wecom-aibot-python-sdk) 的 WSClient 类，
封装了连接管理、心跳、断线重连等能力。

作为独立进程运行：python -m app.wework_bot
"""

from __future__ import annotations

import base64
import hashlib
import json
import logging
import math
import os
import re

import httpx
from aibot import WSClient, WSClientOptions, generate_req_id
from dotenv import load_dotenv

# 加载 .env 文件（独立运行时需要）
load_dotenv()

logger = logging.getLogger("wework_bot")

# 分片上传参数（参考官方文档）
_CHUNK_SIZE = 512 * 1024  # 单个分片最大 512KB（Base64 编码前）


class WeComBotService:
    """企业微信智能机器人服务（WebSocket 长连接模式）"""

    def __init__(self):
        self.bot_id = os.getenv("WEWORK_BOT_ID", "")
        self.bot_secret = os.getenv("WEWORK_BOT_SECRET", "")
        self.edubrain_base_url = os.getenv(
            "EDUBRAIN_BASE_URL", "http://localhost:8765"
        )
        self.search_limit = int(os.getenv("WEWORK_BOT_SEARCH_LIMIT", "3"))
        self.enabled = os.getenv("WEWORK_BOT_ENABLED", "false").lower() == "true"
        self.ws_client: WSClient | None = None

    def start(self):
        """启动机器人服务"""
        if not self.enabled:
            logger.info("企业微信机器人未启用 (WEWORK_BOT_ENABLED=false)")
            return

        if not self.bot_id or not self.bot_secret:
            logger.error("WEWORK_BOT_ID 或 WEWORK_BOT_SECRET 未配置")
            return

        self.ws_client = WSClient(
            WSClientOptions(
                bot_id=self.bot_id,
                secret=self.bot_secret,
            )
        )

        # 注册事件
        self.ws_client.on("authenticated", self._on_authenticated)
        self.ws_client.on("message.text", self._on_text_message)
        self.ws_client.on("event.enter_chat", self._on_enter_chat)
        self.ws_client.on("disconnected", self._on_disconnected)
        self.ws_client.on("error", self._on_error)

        logger.info("启动企业微信智能机器人...")
        self.ws_client.run()

    # ──────────────────────────── 事件处理 ────────────────────────────

    def _on_authenticated(self):
        """认证成功回调"""
        logger.info("企业微信机器人认证成功")

    async def _on_enter_chat(self, frame: dict):
        """进入会话事件 -> 回复欢迎语"""
        await self.ws_client.reply_welcome(
            frame,
            {
                "msgtype": "text",
                "text": {
                    "content": "你好！我是知识库图片搜索助手，发送关键词即可搜索相关图片。"
                },
            },
        )

    async def _on_text_message(self, frame: dict):
        """收到文本消息 -> 搜索图片 -> 流式回复文本 + 逐张发送图片"""
        body = frame.get("body", {})
        text_content = body.get("text", {}).get("content", "")

        # 去掉 @机器人 的前缀
        text_content = re.sub(r"@.*?\s*", "", text_content).strip()

        if not text_content:
            stream_id = generate_req_id("stream")
            await self.ws_client.reply_stream(
                frame, stream_id, "请输入搜索关键词", True
            )
            return

        # 清理口语化前缀，提取关键词
        keyword = re.sub(
            r"(帮我|请|找|搜|搜索|查|查一下|来|给我|要|想要|的?图|图片|相关|关于|一下)",
            "", text_content,
        ).strip()
        if not keyword:
            keyword = text_content  # 兜底

        stream_id = generate_req_id("stream")

        # 立即回复：开始搜索
        await self.ws_client.reply_stream(
            frame,
            stream_id,
            f"\U0001f50d 正在搜索「{keyword}」...",
            False,
        )

        # 调用 EduBrain 搜索
        try:
            results = await self._search_images(keyword)

            if not results:
                await self.ws_client.reply_stream(
                    frame,
                    stream_id,
                    f"未找到与「{text_content}」相关的图片",
                    True,
                )
                return

            # 构建文本回复内容
            content_parts = [f"找到 {len(results)} 张相关图片："]
            for i, r in enumerate(results, 1):
                preview = r.get("ocr_text_preview", "")[:80]
                content_parts.append(f"\n{i}. {preview}...")

            content = "\n".join(content_parts)

            # 完成流式消息（仅文本，不使用 msg_item）
            await self.ws_client.reply_stream(
                frame,
                stream_id,
                content,
                True,
            )

            # 逐张上传临时素材并发送图片消息
            for i, r in enumerate(results, 1):
                b64 = r.get("image_base64", "")
                if not b64:
                    continue

                try:
                    # 解码 base64 获取原始图片数据
                    image_data = base64.b64decode(b64)
                    file_name = r.get("image_url", f"image_{i}.jpg")

                    # 上传临时素材（分片上传）
                    media_id = await self._upload_temp_media(
                        image_data, file_name
                    )
                    if not media_id:
                        logger.warning("第 %d 张图片上传失败，跳过", i)
                        continue

                    # 发送图片消息（通过 aibot_respond_msg）
                    await self.ws_client.reply(
                        frame,
                        {
                            "msgtype": "image",
                            "image": {"media_id": media_id},
                        },
                    )
                    logger.info("已发送第 %d 张图片 (media_id: %s)", i, media_id)
                except Exception as e:
                    logger.error("发送第 %d 张图片失败: %s", i, e, exc_info=True)

        except Exception as e:
            logger.error("搜索失败: %s", e, exc_info=True)
            await self.ws_client.reply_stream(
                frame,
                stream_id,
                f"搜索出错：{e}",
                True,
            )

    def _on_disconnected(self, reason: str):
        """连接断开回调"""
        logger.warning("企业微信连接断开: %s", reason)

    def _on_error(self, error):
        """错误回调"""
        logger.error("企业微信机器人错误: %s", error)

    # ──────────────────────── 上传临时素材 ────────────────────────

    async def _upload_temp_media(
        self, file_data: bytes, filename: str
    ) -> str | None:
        """
        通过 WebSocket 通道上传临时素材（三步分片上传）。

        流程：aibot_upload_media_init → aibot_upload_media_chunk × N → aibot_upload_media_finish

        Args:
            file_data: 文件原始二进制数据
            filename: 文件名

        Returns:
            成功返回 media_id，失败返回 None
        """
        total_size = len(file_data)
        total_chunks = math.ceil(total_size / _CHUNK_SIZE)
        md5 = hashlib.md5(file_data).hexdigest()
        ws = self.ws_client._ws_manager

        # 第一步：上传初始化
        init_req_id = generate_req_id("upload")
        try:
            init_result = await ws.send_reply(
                init_req_id,
                {
                    "type": "image",
                    "filename": filename,
                    "total_size": total_size,
                    "total_chunks": total_chunks,
                    "md5": md5,
                },
                cmd="aibot_upload_media_init",
            )
        except Exception as e:
            logger.error("上传初始化失败: %s", e)
            return None

        upload_id = init_result.get("body", {}).get("upload_id", "")
        if not upload_id:
            logger.error("上传初始化未返回 upload_id: %s", init_result)
            return None

        logger.info(
            "上传初始化成功: upload_id=%s, total_size=%d, chunks=%d",
            upload_id, total_size, total_chunks,
        )

        # 第二步：分片上传
        for chunk_idx in range(total_chunks):
            start = chunk_idx * _CHUNK_SIZE
            end = min(start + _CHUNK_SIZE, total_size)
            chunk_data = file_data[start:end]
            chunk_b64 = base64.b64encode(chunk_data).decode("utf-8")

            chunk_req_id = generate_req_id("upload_chunk")
            try:
                await ws.send_reply(
                    chunk_req_id,
                    {
                        "upload_id": upload_id,
                        "chunk_index": chunk_idx,  # 分片索引从 0 开始
                        "base64_data": chunk_b64,
                    },
                    cmd="aibot_upload_media_chunk",
                )
            except Exception as e:
                logger.error(
                    "分片上传失败 (chunk %d/%d): %s",
                    chunk_idx + 1, total_chunks, e,
                )
                return None

        logger.info("所有分片上传完成: upload_id=%s", upload_id)

        # 第三步：完成上传
        finish_req_id = generate_req_id("upload_finish")
        try:
            finish_result = await ws.send_reply(
                finish_req_id,
                {"upload_id": upload_id},
                cmd="aibot_upload_media_finish",
            )
        except Exception as e:
            logger.error("完成上传失败: %s", e)
            return None

        media_id = finish_result.get("body", {}).get("media_id", "")
        if not media_id:
            logger.error("完成上传未返回 media_id: %s", finish_result)
            return None

        logger.info("上传临时素材成功: media_id=%s", media_id)
        return media_id

    # ──────────────────────────── 搜索调用 ────────────────────────────

    async def _get_search_limit(self) -> int:
        """从后端 API 实时获取当前 search_limit 配置"""
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                resp = await client.get(f"{self.edubrain_base_url}/api/v1/settings")
                if resp.status_code == 200:
                    data = resp.json()
                    return data.get("search_limit", self.search_limit)
        except Exception as e:
            logger.warning("获取 search_limit 失败，使用默认值: %s", e)
        return self.search_limit

    async def _search_images(self, keyword: str) -> list:
        """
        调用 EduBrain 搜索接口（SSE 流式），收集所有匹配结果。

        Args:
            keyword: 搜索关键词

        Returns:
            匹配结果列表
        """
        results = []
        url = f"{self.edubrain_base_url}/api/v1/images/search"
        limit = await self._get_search_limit()

        async with httpx.AsyncClient(timeout=300.0) as client:
            async with client.stream(
                "GET",
                url,
                params={
                    "keyword": keyword,
                    "limit": limit,
                    "sort": "time_desc",
                },
            ) as resp:
                async for line in resp.aiter_lines():
                    if not line.startswith("data: "):
                        continue
                    try:
                        data = json.loads(line[6:])
                        if data.get("type") == "result":
                            results.append(data)
                        elif data.get("type") == "done":
                            break
                    except json.JSONDecodeError:
                        continue

        return results


def run_wework_bot():
    """入口函数"""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    )
    service = WeComBotService()
    service.start()


if __name__ == "__main__":
    run_wework_bot()