"""
企业微信智能机器人服务（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 tempfile
import time
import asyncio
from datetime import datetime

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._register_events()
        logger.info("启动企业微信智能机器人...")
        self.ws_client.run()

    async def start_async(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._register_events()
        logger.info("启动企业微信智能机器人（异步模式）...")
        await self.ws_client.connect()
        # 永远等待，保持连接
        await asyncio.Future()

    def _register_events(self):
        """注册所有事件处理器"""
        self.ws_client.on("authenticated", self._on_authenticated)
        self.ws_client.on("message.text", self._on_text_message)
        self.ws_client.on("message.image", self._on_image_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)

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

    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": "你好！我是知识库助手，支持以下功能：\n\n1. 发送关键词 → 搜索相关图片\n2. 发送图片 → 自动识别并录入知识库"
                },
            },
        )

    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)

    async def _on_image_message(self, frame: dict):
        """收到图片消息 -> 下载解密 -> 调用OCR识别 -> 回复结果"""
        body = frame.get("body", {})
        image_info = body.get("image", {})
        image_url = image_info.get("url", "")
        aes_key = image_info.get("aeskey", "")

        if not image_url:
            logger.warning("收到图片消息但无下载地址")
            return

        stream_id = generate_req_id("stream")
        await self.ws_client.reply_stream(
            frame, stream_id, "收到图片，正在识别...", False
        )

        try:
            # 1. 下载并解密图片
            image_data, original_filename = await self.ws_client.download_file(
                image_url, aes_key
            )
            if not image_data:
                await self.ws_client.reply_stream(
                    frame, stream_id, "图片下载失败，请重试", True
                )
                return

            logger.info(
                "图片下载成功: filename=%s, size=%d bytes",
                original_filename, len(image_data),
            )

            # 2. 生成保存文件名（避免重名）
            ext = os.path.splitext(original_filename or "image.jpg")[1] or ".jpg"
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            save_filename = f"wecom_{timestamp}{ext}"

            # 3. 调用后端 recognize-direct 接口
            async with httpx.AsyncClient(timeout=120.0) as client:
                resp = await client.post(
                    f"{self.edubrain_base_url}/api/v1/images/recognize-direct",
                    files={"file": (save_filename, image_data, "image/jpeg")},
                )

            if resp.status_code != 200:
                logger.error("OCR 识别失败: status=%d, body=%s", resp.status_code, resp.text)
                await self.ws_client.reply_stream(
                    frame, stream_id, "图片识别失败，请稍后重试", True
                )
                return

            result = resp.json()

            # 4. 根据识别结果回复
            status = result.get("status", "")

            if status == "duplicate":
                dup_id = result.get("duplicate_of", "?")
                await self.ws_client.reply_stream(
                    frame,
                    stream_id,
                    f"该图片内容已存在于知识库中（ID: {dup_id}），无需重复录入。",
                    True,
                )
            elif status == "completed":
                ocr_text = result.get("ocr_text", "")
                preview = ocr_text[:200].replace("\n", " ") if ocr_text else "无识别内容"
                record_id = result.get("id", "?")
                await self.ws_client.reply_stream(
                    frame,
                    stream_id,
                    f"✅ 图片已录入知识库（ID: {record_id}）\n\n识别内容预览：\n{preview}{'...' if len(ocr_text) > 200 else ''}",
                    True,
                )
            elif status == "failed":
                error_msg = result.get("message", result.get("detail", "未知错误"))
                await self.ws_client.reply_stream(
                    frame,
                    stream_id,
                    f"图片识别失败：{error_msg}",
                    True,
                )
            else:
                await self.ws_client.reply_stream(
                    frame,
                    stream_id,
                    f"图片处理结果：{json.dumps(result, ensure_ascii=False)}",
                    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_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()