# -*- coding: utf-8 -*-
"""
STDIO MCP Server 实现
通过标准输入输出进行 MCP 协议通信
"""

import sys
import json
import logging
from typing import Dict, List, Any, Optional
import pickle
from pathlib import Path

logger = logging.getLogger(__name__)


class StdioMCPServer:
    """STDIO MCP 协议服务器"""

    def __init__(self):
        """初始化服务器"""
        self.tools: Dict[str, Dict[str, Any]] = {}
        self.resources: Dict[str, Dict[str, Any]] = {}
        self.prompts: Dict[str, Dict[str, Any]] = {}
        self._initialized = False

        # 加载 ML 模型
        self._load_model()

        # 注册核心工具
        self._register_tools()

    def _load_model(self):
        """加载 ML 模型"""
        try:
            # 获取包安装目录
            package_dir = Path(__file__).parent.parent
            model_path = package_dir / "models" / "lightweight_model_all_data.pkl"

            if not model_path.exists():
                # 如果模型不在包目录，尝试从相对路径加载
                model_path = Path(__file__).parent.parent.parent.parent / "tests" / "ml_output" / "lightweight_model_all_data.pkl"

            if model_path.exists():
                with open(model_path, 'rb') as f:
                    model_data = pickle.load(f)
                self.vectorizer = model_data['vectorizer']
                self.model = model_data['model']
                self.reverse_label_map = model_data['reverse_label_map']
                logger.info(f"模型加载成功: {model_path}")
            else:
                logger.warning(f"模型文件不存在: {model_path}")
                self.vectorizer = None
                self.model = None
                self.reverse_label_map = None
        except Exception as e:
            logger.error(f"模型加载失败: {str(e)}")
            self.vectorizer = None
            self.model = None
            self.reverse_label_map = None

    def _register_tools(self):
        """注册核心工具"""
        self.tools["recognize_scenario"] = {
            "name": "recognize_scenario",
            "description": "识别对话场景，支持职场、爱情、家庭、社交四种场景的识别，准确率95%+。输入一段对话文本，返回场景类型和置信度。",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "需要识别场景的对话文本"
                    }
                },
                "required": ["text"]
            }
        }

        self.tools["analyze_participants"] = {
            "name": "analyze_participants",
            "description": "分析对话参与者，识别对话中涉及的人数和角色关系",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "需要分析的对话文本"
                    }
                },
                "required": ["text"]
            }
        }

        self.tools["analyze_relationship"] = {
            "name": "analyze_relationship",
            "description": "分析对话者之间的关系，判断是平等关系、上下级关系还是亲密关系",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "需要分析的对话文本"
                    }
                },
                "required": ["text"]
            }
        }

        self.tools["analyze_attitude"] = {
            "name": "analyze_attitude",
            "description": "分析对话态度，识别说话人的情绪和态度倾向",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "需要分析的对话文本"
                    }
                },
                "required": ["text"]
            }
        }

        logger.info(f"已注册 {len(self.tools)} 个工具")

    def _recognize_scenario(self, text: str) -> Dict[str, Any]:
        """识别场景"""
        if not self.model or not self.vectorizer:
            return {
                "scenario": "unknown",
                "confidence": 0.0,
                "error": "模型未加载"
            }

        try:
            features = self.vectorizer.transform([text])
            pred_label = self.model.predict(features)[0]
            predicted = self.reverse_label_map[pred_label]

            # 获取概率（如果支持）
            confidence = 0.95
            if hasattr(self.model, 'predict_proba'):
                proba = self.model.predict_proba(features)[0]
                confidence = float(proba[pred_label])

            return {
                "scenario": predicted,
                "confidence": round(confidence, 4)
            }
        except Exception as e:
            logger.error(f"场景识别失败: {str(e)}")
            return {
                "scenario": "unknown",
                "confidence": 0.0,
                "error": str(e)
            }

    def _analyze_participants(self, text: str) -> Dict[str, Any]:
        """分析参与者"""
        # 简单的启发式分析
        lines = text.split('\n')
        speakers = set()

        for line in lines:
            line = line.strip()
            if ':' in line:
                speaker = line.split(':')[0].strip()
                if speaker and len(speaker) < 20:  # 排除过长的"说话人"
                    speakers.add(speaker)

        return {
            "participant_count": len(speakers),
            "participants": list(speakers)
        }

    def _analyze_relationship(self, text: str) -> Dict[str, Any]:
        """分析关系"""
        relationship = "unknown"
        confidence = 0.5

        # 简单的关键词分析
        keywords = {
            "平等": ["同事", "同学", "朋友", "伙伴"],
            "上下级": ["老板", "经理", "领导", "下属", "员工", "主管"],
            "亲密": ["亲爱的", "宝贝", "老公", "老婆", "爸", "妈", "孩子"]
        }

        text_lower = text.lower()
        scores = {}

        for rel, words in keywords.items():
            score = sum(1 for word in words if word in text_lower)
            if score > 0:
                scores[rel] = score

        if scores:
            relationship = max(scores, key=scores.get)
            confidence = min(0.9, 0.5 + scores[relationship] * 0.1)

        return {
            "relationship": relationship,
            "confidence": round(confidence, 4)
        }

    def _analyze_attitude(self, text: str) -> Dict[str, Any]:
        """分析态度"""
        attitude = "neutral"
        confidence = 0.5

        # 简单的情感分析
        positive_words = ["好", "棒", "优秀", "赞", "喜欢", "爱", "感谢"]
        negative_words = ["差", "坏", "讨厌", "恨", "不满", "生气", "愤怒"]

        pos_count = sum(1 for word in positive_words if word in text)
        neg_count = sum(1 for word in negative_words if word in text)

        if pos_count > neg_count:
            attitude = "positive"
            confidence = min(0.9, 0.5 + pos_count * 0.1)
        elif neg_count > pos_count:
            attitude = "negative"
            confidence = min(0.9, 0.5 + neg_count * 0.1)

        return {
            "attitude": attitude,
            "confidence": round(confidence, 4)
        }

    async def call_tool(self, name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """调用工具"""
        if name == "recognize_scenario":
            return self._recognize_scenario(arguments.get("text", ""))
        elif name == "analyze_participants":
            return self._analyze_participants(arguments.get("text", ""))
        elif name == "analyze_relationship":
            return self._analyze_relationship(arguments.get("text", ""))
        elif name == "analyze_attitude":
            return self._analyze_attitude(arguments.get("text", ""))
        else:
            raise ValueError(f"未知工具: {name}")

    def list_tools(self) -> List[Dict[str, Any]]:
        """列出所有工具（MCP 协议核心方法）"""
        return [
            {
                "name": tool["name"],
                "description": tool["description"],
                "inputSchema": tool["inputSchema"]
            }
            for tool in self.tools.values()
        ]

    def handle_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """处理 MCP 请求"""
        method = request.get("method")
        params = request.get("params", {})
        req_id = request.get("id")

        response = {
            "jsonrpc": "2.0",
            "id": req_id
        }

        try:
            if method == "initialize":
                self._initialized = True
                response["result"] = {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {
                        "tools": {},
                        "resources": {}
                    },
                    "serverInfo": {
                        "name": "急语箱MCP服务器",
                        "version": "1.2.1"
                    }
                }

            elif method == "tools/list":
                if not self._initialized:
                    raise ValueError("服务器未初始化")
                response["result"] = {
                    "tools": self.list_tools()
                }

            elif method == "tools/call":
                if not self._initialized:
                    raise ValueError("服务器未初始化")
                tool_name = params.get("name")
                arguments = params.get("arguments", {})
                result = self.call_tool(tool_name, arguments)
                response["result"] = {
                    "content": [
                        {
                            "type": "text",
                            "text": json.dumps(result, ensure_ascii=False, indent=2)
                        }
                    ]
                }

            elif method == "resources/list":
                response["result"] = {
                    "resources": []
                }

            elif method == "resources/read":
                # 当前没有实现资源，返回空结果
                uri = params.get("uri")
                response["result"] = {
                    "contents": []
                }

            elif method == "prompts/list":
                # 返回提示词模板列表（当前为空）
                response["result"] = {
                    "prompts": []
                }

            elif method == "prompts/get":
                # 获取特定的提示词模板（当前未实现）
                name = params.get("name")
                arguments = params.get("arguments", {})
                response["error"] = {
                    "code": -32602,
                    "message": f"Prompt not found: {name}"
                }

            elif method == "ping":
                response["result"] = {}

            else:
                response["error"] = {
                    "code": -32601,
                    "message": f"未知方法: {method}"
                }

        except Exception as e:
            logger.error(f"请求处理失败: {str(e)}", exc_info=True)
            response["error"] = {
                "code": -32603,
                "message": str(e)
            }

        return response

    def run(self):
        """运行 STDIO 服务器"""
        logger.info("急语箱 MCP Server (STDIO) 启动")

        try:
            for line in sys.stdin:
                line = line.strip()
                if not line:
                    continue

                try:
                    request = json.loads(line)
                    response = self.handle_request(request)
                    print(json.dumps(response, ensure_ascii=False))
                    sys.stdout.flush()
                except json.JSONDecodeError as e:
                    error_response = {
                        "jsonrpc": "2.0",
                        "id": None,
                        "error": {
                            "code": -32700,
                            "message": f"JSON 解析错误: {str(e)}"
                        }
                    }
                    print(json.dumps(error_response, ensure_ascii=False))
                    sys.stdout.flush()
                except Exception as e:
                    logger.error(f"处理请求失败: {str(e)}", exc_info=True)

        except KeyboardInterrupt:
            logger.info("服务器已停止")
        except Exception as e:
            logger.error(f"服务器错误: {str(e)}", exc_info=True)


def main():
    """主入口"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

    server = StdioMCPServer()
    server.run()


if __name__ == "__main__":
    main()
