#!/usr/bin/env python3
"""
Funstat BOT MCP Server

基于 Telethon 的 MCP 服务器，用于与 @openaiw_bot 交互
提供搜索、查询、统计等功能
"""

import asyncio
import json
import logging
import os
import time
from typing import Any, Dict, List, Optional
from datetime import datetime, timedelta
from collections import deque

from mcp.server import Server
from mcp.types import (
    Resource,
    Tool,
    TextContent,
    ImageContent,
    EmbeddedResource,
)
from pydantic import AnyUrl
from telethon import TelegramClient
from telethon.tl.types import Message

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("funstat_mcp")

# 配置
API_ID = 24660516
API_HASH = "eae564578880a59c9963916ff1bbbd3a"
# Session 文件路径 - 使用独立的安全目录，防止被意外删除
SESSION_PATH = os.path.expanduser("~/telegram_sessions/funstat_bot")
BOT_USERNAME = "@openaiw_bot"

# 速率限制配置
RATE_LIMIT_PER_SECOND = 18  # 每秒最多18个请求
RATE_LIMIT_WINDOW = 1.0  # 1秒时间窗口

# 缓存配置
CACHE_TTL = 3600  # 缓存1小时


class RateLimiter:
    """速率限制器"""

    def __init__(self, max_requests: int, time_window: float):
        self.max_requests = max_requests
        self.time_window = time_window
        self.requests = deque()

    async def acquire(self):
        """获取请求许可，如果超过限制则等待"""
        now = time.time()

        # 移除超出时间窗口的请求记录
        while self.requests and self.requests[0] < now - self.time_window:
            self.requests.popleft()

        # 如果达到限制，等待
        if len(self.requests) >= self.max_requests:
            sleep_time = self.requests[0] + self.time_window - now
            if sleep_time > 0:
                logger.info(f"速率限制: 等待 {sleep_time:.2f} 秒")
                await asyncio.sleep(sleep_time)
                return await self.acquire()  # 递归重试

        # 记录请求时间
        self.requests.append(now)


class ResponseCache:
    """响应缓存"""

    def __init__(self, ttl: int = CACHE_TTL):
        self.cache: Dict[str, tuple[Any, float]] = {}
        self.ttl = ttl

    def get(self, key: str) -> Optional[Any]:
        """获取缓存"""
        if key in self.cache:
            value, timestamp = self.cache[key]
            if time.time() - timestamp < self.ttl:
                logger.info(f"缓存命中: {key}")
                return value
            else:
                # 过期，删除
                del self.cache[key]
        return None

    def set(self, key: str, value: Any):
        """设置缓存"""
        self.cache[key] = (value, time.time())
        logger.info(f"缓存保存: {key}")

    def clear_expired(self):
        """清理过期缓存"""
        now = time.time()
        expired_keys = [
            key for key, (_, timestamp) in self.cache.items()
            if now - timestamp >= self.ttl
        ]
        for key in expired_keys:
            del self.cache[key]
        if expired_keys:
            logger.info(f"清理了 {len(expired_keys)} 个过期缓存")


class FunstatMCPServer:
    """Funstat MCP 服务器"""

    def __init__(self):
        self.server = Server("funstat-mcp")
        self.client: Optional[TelegramClient] = None
        self.bot_entity = None
        self.rate_limiter = RateLimiter(RATE_LIMIT_PER_SECOND, RATE_LIMIT_WINDOW)
        self.cache = ResponseCache()

        # 注册处理器
        self.server.list_tools()(self.list_tools)
        self.server.call_tool()(self.call_tool)

    async def initialize(self):
        """初始化 Telegram 客户端"""
        logger.info("初始化 Telegram 客户端...")

        # 检查 session 文件
        session_file = f"{SESSION_PATH}.session"
        if not os.path.exists(session_file):
            raise FileNotFoundError(
                f"Session 文件不存在: {session_file}\n"
                f"请先运行 create_session.py 创建 session 文件\n"
                f"或者将现有 session 文件复制到: ~/telegram_sessions/"
            )

        logger.info(f"使用 Session 文件: {session_file}")

        # 创建客户端
        self.client = TelegramClient(SESSION_PATH, API_ID, API_HASH)
        await self.client.start()

        # 获取 bot 实体
        logger.info(f"连接到 {BOT_USERNAME}...")
        self.bot_entity = await self.client.get_entity(BOT_USERNAME)
        logger.info(f"✅ 已连接到: {self.bot_entity.first_name}")

        # 获取当前用户信息
        me = await self.client.get_me()
        logger.info(f"✅ 当前账号: @{me.username} (ID: {me.id})")

    async def send_command_and_wait(
        self,
        command: str,
        timeout: int = 10,
        use_cache: bool = True
    ) -> str:
        """发送命令到 BOT 并等待响应"""

        # 检查缓存
        cache_key = f"cmd:{command}"
        if use_cache:
            cached = self.cache.get(cache_key)
            if cached:
                return cached

        # 速率限制
        await self.rate_limiter.acquire()

        logger.info(f"📤 发送命令: {command}")

        # 记录发送前的最新消息 ID
        last_message_id = 0
        async for message in self.client.iter_messages(self.bot_entity, limit=1):
            last_message_id = message.id
            break

        # 发送消息
        send_time = datetime.now()
        await self.client.send_message(self.bot_entity, command)

        # 等待响应（稍等一下让 BOT 有时间响应）
        await asyncio.sleep(1.5)

        # 获取新消息
        start_time = time.time()
        while time.time() - start_time < timeout:
            # 获取最新消息
            async for message in self.client.iter_messages(self.bot_entity, limit=5):
                # 检查是否是新消息
                if message.id > last_message_id:
                    # 检查是否是 BOT 的消息
                    if not message.out and message.text:
                        response_text = message.text
                        logger.info(f"✅ 收到响应 ({len(response_text)} 字符)")

                        # 保存到缓存
                        if use_cache:
                            self.cache.set(cache_key, response_text)

                        return response_text

            # 继续等待
            await asyncio.sleep(0.5)

        raise TimeoutError(f"等待 BOT 响应超时 ({timeout}秒)")

    async def list_tools(self) -> List[Tool]:
        """列出所有可用工具"""
        return [
            Tool(
                name="funstat_search",
                description="搜索 Telegram 群组、频道。支持关键词搜索，返回相关的群组列表",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "搜索关键词，例如: 'python', '区块链', 'AI'"
                        }
                    },
                    "required": ["query"]
                }
            ),
            Tool(
                name="funstat_topchat",
                description="获取热门群组/频道列表，按成员数或活跃度排序",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "category": {
                            "type": "string",
                            "description": "分类筛选（可选），例如: 'tech', 'crypto', 'news'"
                        }
                    }
                }
            ),
            Tool(
                name="funstat_text",
                description="通过消息文本搜索，查找包含特定文本的消息和来源群组",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "text": {
                            "type": "string",
                            "description": "要搜索的文本内容"
                        }
                    },
                    "required": ["text"]
                }
            ),
            Tool(
                name="funstat_human",
                description="通过姓名搜索用户，查找 Telegram 用户信息",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "name": {
                            "type": "string",
                            "description": "用户姓名"
                        }
                    },
                    "required": ["name"]
                }
            ),
            Tool(
                name="funstat_user_info",
                description="查询用户详细信息，支持通过用户名、用户ID、联系人等方式查询",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "identifier": {
                            "type": "string",
                            "description": "用户标识: 用户名(@username)、用户ID、或手机号"
                        }
                    },
                    "required": ["identifier"]
                }
            ),
            Tool(
                name="funstat_balance",
                description="查询当前账号的积分余额和使用统计",
                inputSchema={
                    "type": "object",
                    "properties": {}
                }
            ),
            Tool(
                name="funstat_menu",
                description="显示 funstat BOT 的主菜单和所有可用功能",
                inputSchema={
                    "type": "object",
                    "properties": {}
                }
            ),
            Tool(
                name="funstat_start",
                description="获取 funstat BOT 的欢迎信息和使用说明",
                inputSchema={
                    "type": "object",
                    "properties": {}
                }
            )
        ]

    async def call_tool(self, name: str, arguments: Dict[str, Any]) -> List[TextContent]:
        """调用工具"""
        logger.info(f"🔧 调用工具: {name} with {arguments}")

        try:
            if name == "funstat_search":
                query = arguments["query"]
                response = await self.send_command_and_wait(f"/search {query}")
                return [TextContent(type="text", text=response)]

            elif name == "funstat_topchat":
                category = arguments.get("category", "")
                if category:
                    response = await self.send_command_and_wait(f"/topchat {category}")
                else:
                    response = await self.send_command_and_wait("/topchat")
                return [TextContent(type="text", text=response)]

            elif name == "funstat_text":
                text = arguments["text"]
                response = await self.send_command_and_wait(f"/text {text}")
                return [TextContent(type="text", text=response)]

            elif name == "funstat_human":
                name_query = arguments["name"]
                response = await self.send_command_and_wait(f"/human {name_query}")
                return [TextContent(type="text", text=response)]

            elif name == "funstat_user_info":
                identifier = arguments["identifier"]
                # 直接发送用户标识（用户名、ID等）
                response = await self.send_command_and_wait(identifier)
                return [TextContent(type="text", text=response)]

            elif name == "funstat_balance":
                response = await self.send_command_and_wait("/balance")
                return [TextContent(type="text", text=response)]

            elif name == "funstat_menu":
                response = await self.send_command_and_wait("/menu")
                return [TextContent(type="text", text=response)]

            elif name == "funstat_start":
                response = await self.send_command_and_wait("/start")
                return [TextContent(type="text", text=response)]

            else:
                raise ValueError(f"未知工具: {name}")

        except Exception as e:
            logger.error(f"❌ 工具调用失败: {e}")
            return [TextContent(
                type="text",
                text=f"❌ 错误: {str(e)}"
            )]

    async def run(self):
        """运行服务器"""
        await self.initialize()

        # 启动定期清理过期缓存的任务
        async def cache_cleanup_task():
            while True:
                await asyncio.sleep(300)  # 每5分钟清理一次
                self.cache.clear_expired()

        asyncio.create_task(cache_cleanup_task())

        logger.info("🚀 Funstat MCP Server 已启动")

        # 运行 MCP 服务器
        from mcp.server.stdio import stdio_server

        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                self.server.create_initialization_options()
            )


async def main():
    """主函数"""
    server = FunstatMCPServer()
    await server.run()


if __name__ == "__main__":
    asyncio.run(main())