#!/usr/bin/env python3
"""
消息服务
提供加载会话存档、解密、保存到数据库的整合功能
"""

import asyncio
import json
import os
import logging
import base64
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any, Optional
import sys

# 添加scripts目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))

from wework_ctypes import WeWorkFinanceSDK, WeWorkSDKError
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.backends import default_backend

logger = logging.getLogger(__name__)


class MessageService:
    """消息服务，提供加载会话存档的整合功能"""

    def __init__(self, corpid: Optional[str] = None, secret: Optional[str] = None):
        """
        初始化消息服务

        Args:
            corpid: 企业ID，如果为None则从环境变量CORPID读取
            secret: 聊天内容存档的Secret，如果为None则从环境变量WEWORK_SECRET读取
        """
        self.corpid = corpid or os.getenv("CORPID", "wwdf65802ca25ec195")
        self.secret = secret or os.getenv(
            "WEWORK_SECRET", "-Ta6WMWxBhfGolWnnlO15nQckj3DRKAowUOdX2fwvzE"
        )

        # SDK实例（延迟初始化）
        self._sdk: Optional[WeWorkFinanceSDK] = None

        # RSA私钥路径
        self.rsa_private_key_path = Path("cert/1_rsa.key")

        logger.info(f"消息服务初始化完成，企业ID: {self.corpid}")

    @property
    def sdk(self) -> WeWorkFinanceSDK:
        """获取SDK实例（延迟初始化）"""
        if self._sdk is None:
            try:
                # 动态库路径
                lib_path = (
                    Path(__file__).parent.parent.parent
                    / "lib"
                    / "sdk"
                    / "libWeWorkFinanceSdk_C.so"
                )
                if not lib_path.exists():
                    raise FileNotFoundError(f"动态库不存在: {lib_path}")

                self._sdk = WeWorkFinanceSDK(self.corpid, self.secret, str(lib_path))
                logger.info("企业微信SDK初始化成功")
            except Exception as e:
                logger.error(f"企业微信SDK初始化失败: {e}")
                raise
        return self._sdk

    def _load_rsa_private_key(self, key_path: str) -> str:
        """
        加载RSA私钥

        Args:
            key_path: 私钥文件路径

        Returns:
            私钥内容字符串
        """
        try:
            with open(key_path, "r") as f:
                key_content = f.read()
            return key_content
        except Exception as e:
            raise Exception(f"加载RSA私钥失败: {e}")

    def _rsa_decrypt_base64(self, encrypted_data_base64: str, private_key: str) -> str:
        """
        使用RSA私钥解密Base64编码的数据

        Args:
            encrypted_data_base64: Base64编码的加密数据
            private_key: RSA私钥字符串

        Returns:
            解密后的字符串
        """
        try:
            # 加载私钥
            private_key_obj = serialization.load_pem_private_key(
                private_key.encode("utf-8"), password=None, backend=default_backend()
            )

            # Base64解码
            encrypted_data = base64.b64decode(encrypted_data_base64)

            # RSA解密 (使用PKCS1v15填充)
            decrypted_data = private_key_obj.decrypt(encrypted_data, padding.PKCS1v15())

            return decrypted_data.decode("utf-8")

        except ImportError:
            raise Exception("需要安装cryptography库: pip install cryptography")
        except Exception as e:
            raise Exception(f"RSA解密失败: {e}")

    def _decrypt_chat_message(
        self, encrypt_random_key: str, encrypt_chat_msg: str, publickey_ver: int = None
    ) -> dict:
        """
        解密聊天消息

        Args:
            encrypt_random_key: 加密的随机密钥（Base64编码）
            encrypt_chat_msg: 加密的聊天消息
            publickey_ver: publickey_ver版本号，用于动态匹配私钥

        Returns:
            解密后的消息字典
        """
        # 动态确定私钥路径
        actual_key_path = self.rsa_private_key_path

        # 如果提供了publickey_ver，则动态匹配
        if publickey_ver is not None:
            dynamic_key_path = f"cert/{publickey_ver}_rsa.key"
            dynamic_path = Path(dynamic_key_path)
            if dynamic_path.exists():
                actual_key_path = dynamic_path

        # 1. 加载RSA私钥
        private_key = self._load_rsa_private_key(str(actual_key_path))

        # 2. 使用RSA私钥解密encrypt_random_key得到encrypt_key
        try:
            encrypt_key = self._rsa_decrypt_base64(encrypt_random_key, private_key)
        except Exception as e:
            # 尝试直接使用encrypt_random_key作为encrypt_key（某些情况下可能已经解密）
            encrypt_key = encrypt_random_key

        # 3. 使用SDK解密聊天消息
        decrypted_msg = self.sdk.decrypt_data(encrypt_key, encrypt_chat_msg)

        # 4. 解析JSON
        try:
            result = json.loads(decrypted_msg)
            return result
        except json.JSONDecodeError:
            # 如果不是JSON，返回原始文本
            return {"raw_message": decrypted_msg}

    async def _decrypt_chat_data(self, chatdata: List[Dict[str, Any]]) -> tuple:
        """
        解密聊天数据

        Args:
            chatdata: 原始聊天数据列表

        Returns:
            (decrypted_chatdata, stats)
            decrypted_chatdata: 解密后的聊天数据列表
            stats: 解密统计信息字典
        """
        decrypted_chatdata = []
        stats = {
            "total_encrypted_messages": 0,
            "successful_decrypts": 0,
            "failed_decrypts": 0,
        }

        # 检查私钥文件是否存在
        if not self.rsa_private_key_path.exists():
            logger.warning(f"RSA私钥文件不存在: {self.rsa_private_key_path}")
            logger.warning("将跳过解密，保存原始加密数据")
            return chatdata, stats

        logger.info(f"使用RSA私钥: {self.rsa_private_key_path.absolute()}")
        logger.info("开始解密聊天数据...")

        for i, msg in enumerate(chatdata):
            encrypt_random_key = msg.get("encrypt_random_key", "")
            encrypt_chat_msg = msg.get("encrypt_chat_msg", "")

            # 检查是否有加密数据
            if not encrypt_random_key or not encrypt_chat_msg:
                # 没有加密数据，直接添加到结果中
                decrypted_chatdata.append(msg)
                continue

            stats["total_encrypted_messages"] += 1

            try:
                # 获取publickey_ver
                publickey_ver = msg.get("publickey_ver")

                # 解密消息
                decrypted_msg = self._decrypt_chat_message(
                    encrypt_random_key,
                    encrypt_chat_msg,
                    publickey_ver,
                )

                # 创建新的消息对象，保留原始字段但移除加密字段
                new_msg = {}
                for key, value in msg.items():
                    if key not in ["encrypt_random_key", "encrypt_chat_msg"]:
                        new_msg[key] = value

                # 将解密后的字段合并到消息中
                if isinstance(decrypted_msg, dict):
                    for key, value in decrypted_msg.items():
                        new_msg[key] = value

                decrypted_chatdata.append(new_msg)
                stats["successful_decrypts"] += 1

                # 显示进度
                if stats["successful_decrypts"] % 100 == 0:
                    logger.info(f"  已解密 {stats['successful_decrypts']} 条消息...")

            except Exception as e:
                stats["failed_decrypts"] += 1
                logger.warning(f"解密消息失败 (seq: {msg.get('seq', 'N/A')}): {e}")
                # 解密失败，跳过此消息（不添加到结果中）
                continue

        return decrypted_chatdata, stats

    async def _save_messages_to_database(
        self, messages: List[Dict[str, Any]]
    ) -> Dict[str, int]:
        """
        将消息保存到数据库

        Args:
            messages: 消息字典列表

        Returns:
            存储统计信息字典
        """
        stats = {
            "total_messages": len(messages),
            "saved": 0,
            "skipped_duplicates": 0,
            "failed": 0,
        }

        if not messages:
            logger.info("没有消息需要存储到数据库")
            return stats

        logger.info(f"开始将 {len(messages)} 条消息存储到数据库...")

        try:
            # 导入数据库相关模块
            from pymongo import AsyncMongoClient
            from beanie import init_beanie
            from app.models import models
            from app.models.message import Message

            # 从环境变量获取数据库配置
            mongodb_url = os.getenv(
                "MONGODB_URL",
                "mongodb://dblink:mgdb4link@localhost:27017/?authSource=admin",
            )
            database_name = os.getenv("DATABASE_NAME", "douba_chat_dev")

            # 初始化数据库连接
            client = AsyncMongoClient(mongodb_url)
            await init_beanie(database=client[database_name], document_models=models)

            logger.info(f"数据库连接成功: {database_name}")

            # 批量处理消息
            batch_size = 100
            for i in range(0, len(messages), batch_size):
                batch = messages[i : i + batch_size]
                batch_msgids = [msg.get("msgid") for msg in batch if msg.get("msgid")]

                # 检查哪些消息已经存在
                existing_messages = await Message.find(
                    {"msgid": {"$in": batch_msgids}}
                ).to_list()
                existing_msgids = {msg.msgid for msg in existing_messages}

                # 过滤掉已存在的消息
                new_messages = []
                for msg in batch:
                    msgid = msg.get("msgid")
                    if not msgid:
                        stats["failed"] += 1
                        continue

                    if msgid in existing_msgids:
                        stats["skipped_duplicates"] += 1
                        continue

                    try:
                        # 创建Message实例
                        message_data = {
                            "msgid": msg.get("msgid"),
                            "action": msg.get("action"),
                            "from_user": msg.get(
                                "from"
                            ),  # 从原始数据的from字段读取，保存为from_user
                            "tolist": msg.get("tolist", []),
                            "roomid": msg.get("roomid"),
                            "text": msg.get("text", {}),
                            "msgtime": msg.get("msgtime"),
                            "msgtype": msg.get("msgtype"),
                            "remark": msg.get("remark"),
                            "seq": msg.get("seq", 1000),
                            "isEnabled": True,
                            "created_at": datetime.now(),
                            "updated_at": datetime.now(),
                        }

                        # 移除None值
                        message_data = {
                            k: v for k, v in message_data.items() if v is not None
                        }

                        message = Message(**message_data)
                        new_messages.append(message)
                    except Exception as e:
                        logger.error(f"创建消息实例失败 (msgid: {msgid}): {e}")
                        stats["failed"] += 1

                # 批量插入新消息
                if new_messages:
                    try:
                        await Message.insert_many(new_messages)
                        stats["saved"] += len(new_messages)
                        logger.info(
                            f"  已存储 {len(new_messages)} 条消息 (进度: {i + len(batch)}/{len(messages)})"
                        )
                    except Exception as e:
                        logger.error(f"批量插入失败: {e}")
                        stats["failed"] += len(new_messages)

            # 关闭数据库连接
            await client.close()

            logger.info(f"数据库存储完成:")
            logger.info(f"  总计消息: {stats['total_messages']}")
            logger.info(f"  成功存储: {stats['saved']}")
            logger.info(f"  跳过重复: {stats['skipped_duplicates']}")
            logger.info(f"  存储失败: {stats['failed']}")

        except Exception as e:
            logger.error(f"数据库存储过程中发生错误: {e}")
            import traceback

            traceback.print_exc()
            stats["failed"] = len(messages)

        return stats

    async def get_messages(
        self,
        page: int = 1,
        page_size: int = 20,
        msgtype: Optional[str] = None,
        roomid: Optional[str] = None,
        from_user: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        查询消息列表，支持分页和过滤

        Args:
            page: 页码（从1开始）
            page_size: 每页数量
            msgtype: 消息类型过滤
            roomid: 群聊ID过滤
            from_user: 发送者过滤

        Returns:
            分页查询结果
        """
        try:
            # 导入数据库相关模块
            from pymongo import AsyncMongoClient
            from beanie import init_beanie
            from app.models import models
            from app.models.message import Message

            # 从环境变量获取数据库配置
            mongodb_url = os.getenv(
                "MONGODB_URL",
                "mongodb://dblink:mgdb4link@localhost:27017/?authSource=admin",
            )
            database_name = os.getenv("DATABASE_NAME", "douba_chat_dev")

            # 初始化数据库连接
            client = AsyncMongoClient(mongodb_url)
            await init_beanie(database=client[database_name], document_models=models)

            # 构建查询条件
            query = {}
            if msgtype:
                query["msgtype"] = msgtype
            if roomid:
                query["roomid"] = roomid
            if from_user:
                query["from_user"] = from_user  # 使用from_user字段进行查询

            # 计算跳过的文档数量
            skip_count = (page - 1) * page_size

            # 执行查询
            messages = await Message.find(
                query,
                skip=skip_count,
                limit=page_size,
                sort=[("msgtime", -1)],  # 按消息时间倒序排列
            ).to_list()

            # 获取总记录数
            total_count = await Message.find(query).count()

            # 关闭数据库连接
            await client.close()

            # 准备响应数据
            return {
                "success": True,
                "data": {
                    "page": page,
                    "page_size": page_size,
                    "total": total_count,
                    "total_pages": (
                        (total_count + page_size - 1) // page_size
                        if page_size > 0
                        else 0
                    ),
                    "messages": messages,
                },
            }

        except Exception as e:
            logger.error(f"查询消息列表失败: {e}")
            import traceback

            traceback.print_exc()
            return {
                "success": False,
                "message": f"查询消息列表失败: {str(e)}",
                "data": {
                    "page": page,
                    "page_size": page_size,
                    "total": 0,
                    "total_pages": 0,
                    "messages": [],
                },
            }

    async def load_messages(self, seq: int = 0, limit: int = 100) -> Dict[str, Any]:
        """
        加载会话存档到数据库
        完整流程：拉取 → 解密 → 保存

        Args:
            seq: 起始seq值，首次使用0
            limit: 拉取数量限制

        Returns:
            执行结果和统计信息
        """
        start_time = datetime.now()

        try:
            # 从环境变量获取代理配置
            proxy = os.getenv("WEWORK_PROXY", "")
            passwd = os.getenv("WEWORK_PROXY_PASSWD", "")
            timeout = int(os.getenv("WEWORK_TIMEOUT", "10"))

            logger.info(f"开始加载会话存档: seq={seq}, limit={limit}")

            # 1. 拉取会话数据
            logger.info("步骤1: 从企业微信拉取会话数据...")
            result_str = self.sdk.get_chat_data(
                seq=seq, limit=limit, proxy=proxy, passwd=passwd, timeout=timeout
            )

            # 解析结果
            result = json.loads(result_str)
            errcode = result.get("errcode", -1)
            errmsg = result.get("errmsg", "")

            if errcode != 0:
                raise Exception(f"拉取会话数据失败: errcode={errcode}, errmsg={errmsg}")

            chatdata = result.get("chatdata", [])
            total_fetched = len(chatdata)
            logger.info(f"成功拉取 {total_fetched} 条聊天记录")

            if total_fetched == 0:
                return {
                    "success": True,
                    "message": "没有新的聊天记录",
                    "data": {
                        "total_fetched": 0,
                        "encrypted_count": 0,
                        "decrypted_success": 0,
                        "decrypted_failed": 0,
                        "saved_count": 0,
                        "duplicates_skipped": 0,
                        "max_seq": seq,
                        "next_seq": seq,
                    },
                    "execution_time": str(datetime.now() - start_time),
                }

            # 2. 解密消息
            logger.info("步骤2: 解密加密消息...")
            decrypted_chatdata, decrypt_stats = await self._decrypt_chat_data(chatdata)

            # 3. 保存到数据库
            logger.info("步骤3: 保存消息到数据库...")
            db_stats = await self._save_messages_to_database(decrypted_chatdata)

            # 计算最大seq值
            max_seq = max(
                (msg.get("seq", 0) for msg in decrypted_chatdata), default=seq
            )

            # 准备响应数据
            execution_time = datetime.now() - start_time

            response_data = {
                "success": True,
                "message": "消息加载完成",
                "data": {
                    "total_fetched": total_fetched,
                    "encrypted_count": decrypt_stats["total_encrypted_messages"],
                    "decrypted_success": decrypt_stats["successful_decrypts"],
                    "decrypted_failed": decrypt_stats["failed_decrypts"],
                    "saved_count": db_stats["saved"],
                    "duplicates_skipped": db_stats["skipped_duplicates"],
                    "failed_to_save": db_stats["failed"],
                    "max_seq": max_seq,
                    "next_seq": max_seq + 1 if total_fetched > 0 else seq,
                },
                "execution_time": str(execution_time),
                "timestamp": datetime.now().isoformat(),
            }

            logger.info(f"消息加载完成，耗时: {execution_time}")
            return response_data

        except json.JSONDecodeError as e:
            error_msg = f"解析JSON数据失败: {e}"
            logger.error(error_msg)
            return {
                "success": False,
                "message": error_msg,
                "execution_time": str(datetime.now() - start_time),
            }

        except Exception as e:
            error_msg = f"加载消息失败: {e}"
            logger.error(error_msg)
            import traceback

            traceback.print_exc()
            return {
                "success": False,
                "message": error_msg,
                "execution_time": str(datetime.now() - start_time),
            }

    def close(self):
        """关闭服务，释放资源"""
        if self._sdk:
            self._sdk.close()
            self._sdk = None
            logger.info("消息服务已关闭")


# 全局服务实例（单例模式）
_message_service: Optional[MessageService] = None


def get_message_service() -> MessageService:
    """获取消息服务实例（单例）"""
    global _message_service
    if _message_service is None:
        _message_service = MessageService()
    return _message_service


def close_message_service():
    """关闭消息服务"""
    global _message_service
    if _message_service:
        _message_service.close()
        _message_service = None