#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
思维模型增强器 - P3优化功能模块
Thinking Model Enhancer - P3 Optimization Module

实现P3优化功能：
- 跨会话持续学习
- 多模型协同
- 高级分析功能
- 智能优化建议
"""

import json
import hashlib
from pathlib import Path
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
from collections import defaultdict
import statistics


@dataclass
class LearningPattern:
    """学习模式"""
    pattern_id: str
    pattern_type: str  # successful_pattern, failed_pattern, improvement_area
    description: str
    conditions: List[str]
    actions: List[str]
    success_rate_impact: float  # 对成功率的影响
    occurrence_count: int
    last_observed: str
    confidence: float  # 0-1
    source_sessions: List[str]
    
    def to_dict(self) -> Dict:
        return {
            "pattern_id": self.pattern_id,
            "pattern_type": self.pattern_type,
            "description": self.description,
            "conditions": self.conditions,
            "actions": self.actions,
            "success_rate_impact": self.success_rate_impact,
            "occurrence_count": self.occurrence_count,
            "last_observed": self.last_observed,
            "confidence": self.confidence,
            "source_sessions": self.source_sessions
        }


@dataclass
class SessionData:
    """会话数据"""
    session_id: str
    start_time: str
    end_time: Optional[str]
    user_input: str
    selected_model: str
    stages_completed: List[str]
    result_summary: str
    success: bool
    user_rating: Optional[int]
    duration_ms: float
    insights: List[str] = field(default_factory=list)
    improvements: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict:
        return {
            "session_id": self.session_id,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "user_input": self.user_input,
            "selected_model": self.selected_model,
            "stages_completed": self.stages_completed,
            "result_summary": self.result_summary,
            "success": self.success,
            "user_rating": self.user_rating,
            "duration_ms": self.duration_ms,
            "insights": self.insights,
            "improvements": self.improvements
        }


@dataclass
class ModelCollaboration:
    """模型协作配置"""
    collaboration_id: str
    name: str
    description: str
    models: List[str]  # 参与的模型
    workflow: List[Dict]  # 工作流程定义
    conditions: List[str]  # 适用条件
    performance_metrics: Dict = field(default_factory=dict)
    use_count: int = 0
    
    def to_dict(self) -> Dict:
        return {
            "collaboration_id": self.collaboration_id,
            "name": self.name,
            "description": self.description,
            "models": self.models,
            "workflow": self.workflow,
            "conditions": self.conditions,
            "performance_metrics": self.performance_metrics,
            "use_count": self.use_count
        }


class ContinuousLearning:
    """跨会话持续学习系统"""
    
    def __init__(self, learning_dir: Optional[str] = None):
        """初始化持续学习系统"""
        if learning_dir is None:
            self.learning_dir = Path.home() / ".claude" / "thinking_models" / "learning"
        else:
            self.learning_dir = Path(learning_dir)
        
        self.learning_dir.mkdir(parents=True, exist_ok=True)
        
        # 数据文件
        self.patterns_file = self.learning_dir / "patterns.json"
        self.sessions_file = self.learning_dir / "sessions.json"
        
        # 加载数据
        self.patterns = self._load_patterns()
        self.sessions = self._load_sessions()
    
    def _load_patterns(self) -> Dict:
        """加载学习模式"""
        if self.patterns_file.exists():
            with open(self.patterns_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        return {
            "version": "1.0",
            "updated_at": datetime.now().isoformat(),
            "patterns": []
        }
    
    def _load_sessions(self) -> Dict:
        """加载会话数据"""
        if self.sessions_file.exists():
            with open(self.sessions_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        return {
            "version": "1.0",
            "updated_at": datetime.now().isoformat(),
            "sessions": []
        }
    
    def _save_patterns(self):
        """保存学习模式"""
        self.patterns["updated_at"] = datetime.now().isoformat()
        with open(self.patterns_file, 'w', encoding='utf-8') as f:
            json.dump(self.patterns, f, ensure_ascii=False, indent=2)
    
    def _save_sessions(self):
        """保存会话数据"""
        self.sessions["updated_at"] = datetime.now().isoformat()
        with open(self.sessions_file, 'w', encoding='utf-8') as f:
            json.dump(self.sessions, f, ensure_ascii=False, indent=2)
    
    def record_session(self, session: Dict) -> str:
        """
        记录一个会话
        
        Args:
            session: 会话数据字典
            
        Returns:
            会话ID
        """
        session_id = session.get("session_id") or f"session_{datetime.now().strftime('%Y%m%d%H%M%S')}_{hashlib.md5(str(datetime.now()).encode()).hexdigest()[:6]}"
        session["session_id"] = session_id
        
        self.sessions["sessions"].append(session)
        self._save_sessions()
        
        # 触发模式学习
        self._learn_from_session(session)
        
        return session_id
    
    def _learn_from_session(self, session: Dict):
        """从会话中学习模式"""
        # 分析成功模式
        if session.get("success", False):
            self._extract_successful_pattern(session)
        
        # 分析改进建议
        improvements = session.get("improvements", [])
        if improvements:
            self._extract_improvement_pattern(session, improvements)
    
    def _extract_successful_pattern(self, session: Dict):
        """提取成功模式"""
        # 简化实现：基于模型和问题类型创建模式
        model = session.get("selected_model", "")
        problem_type = session.get("problem_type", "general")
        
        pattern_desc = f"使用{model}处理{problem_type}类型问题成功"
        
        # 检查是否已存在类似模式
        existing = None
        for p in self.patterns.get("patterns", []):
            if p.get("description") == pattern_desc:
                existing = p
                break
        
        if existing:
            existing["occurrence_count"] += 1
            existing["last_observed"] = datetime.now().isoformat()
            existing["confidence"] = min(1.0, existing["occurrence_count"] / 10)
        else:
            new_pattern = {
                "pattern_id": f"pattern_{hashlib.md5((pattern_desc + datetime.now().isoformat()).encode()).hexdigest()[:8]}",
                "pattern_type": "successful_pattern",
                "description": pattern_desc,
                "conditions": [problem_type],
                "actions": [model],
                "success_rate_impact": 0.05,
                "occurrence_count": 1,
                "last_observed": datetime.now().isoformat(),
                "confidence": 0.3,
                "source_sessions": [session.get("session_id", "")]
            }
            self.patterns["patterns"].append(new_pattern)
        
        self._save_patterns()
    
    def _extract_improvement_pattern(self, session: Dict, improvements: List[str]):
        """提取改进模式"""
        model = session.get("selected_model", "")
        
        pattern_desc = f"{model}的改进建议: {', '.join(improvements[:2])}"
        
        new_pattern = {
            "pattern_id": f"pattern_{hashlib.md5((pattern_desc + datetime.now().isoformat()).encode()).hexdigest()[:8]}",
            "pattern_type": "improvement_area",
            "description": pattern_desc,
            "conditions": [session.get("problem_type", "")],
            "actions": improvements,
            "success_rate_impact": 0.02,
            "occurrence_count": 1,
            "last_observed": datetime.now().isoformat(),
            "confidence": 0.2,
            "source_sessions": [session.get("session_id", "")]
        }
        
        self.patterns["patterns"].append(new_pattern)
        self._save_patterns()
    
    def get_patterns_by_type(self, pattern_type: str) -> List[Dict]:
        """按类型获取模式"""
        return [p for p in self.patterns.get("patterns", []) if p.get("pattern_type") == pattern_type]
    
    def get_relevant_patterns(self, problem_type: str, model: str) -> List[Dict]:
        """获取相关模式"""
        relevant = []
        for p in self.patterns.get("patterns", []):
            # 检查条件匹配
            conditions = p.get("conditions", [])
            if problem_type in conditions or not conditions:
                # 检查动作匹配
                actions = p.get("actions", [])
                if model in actions or not actions:
                    relevant.append(p)
        
        # 按置信度排序
        relevant.sort(key=lambda x: x.get("confidence", 0), reverse=True)
        return relevant
    
    def get_learning_summary(self) -> Dict:
        """获取学习摘要"""
        patterns = self.patterns.get("patterns", [])
        sessions = self.sessions.get("sessions", [])
        
        successful_patterns = [p for p in patterns if p.get("pattern_type") == "successful_pattern"]
        improvement_patterns = [p for p in patterns if p.get("pattern_type") == "improvement_area"]
        
        return {
            "total_patterns": len(patterns),
            "successful_patterns": len(successful_patterns),
            "improvement_patterns": len(improvement_patterns),
            "total_sessions": len(sessions),
            "avg_session_duration": statistics.mean([s.get("duration_ms", 0) for s in sessions]) if sessions else 0,
            "top_patterns": sorted(patterns, key=lambda x: x.get("confidence", 0), reverse=True)[:5]
        }


class ModelCollaborator:
    """多模型协同管理器"""
    
    def __init__(self, collab_dir: Optional[str] = None):
        """初始化多模型协同管理器"""
        if collab_dir is None:
            self.collab_dir = Path.home() / ".claude" / "thinking_models" / "collaborations"
        else:
            self.collab_dir = Path(collab_dir)
        
        self.collab_dir.mkdir(parents=True, exist_ok=True)
        self.collaborations_file = self.collab_dir / "collaborations.json"
        
        # 加载或创建协作配置
        self.collaborations = self._load_or_create_collaborations()
    
    def _load_or_create_collaborations(self) -> Dict:
        """加载或创建协作配置"""
        if self.collaborations_file.exists():
            with open(self.collaborations_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        
        # 默认协作配置
        default_collabs = {
            "version": "1.0",
            "updated_at": datetime.now().isoformat(),
            "collaborations": [
                {
                    "collaboration_id": "research_diagnostic",
                    "name": "研究-诊断协同",
                    "description": "先进行研究分析，再进行诊断修复",
                    "models": ["research_mode", "diagnostic_mode"],
                    "workflow": [
                        {"order": 1, "model": "research_mode", "action": "分析问题背景"},
                        {"order": 2, "model": "research_mode", "action": "收集解决方案"},
                        {"order": 3, "model": "diagnostic_mode", "action": "诊断具体问题"},
                        {"order": 4, "model": "diagnostic_mode", "action": "制定修复方案"}
                    ],
                    "conditions": ["复杂问题", "需要全面分析"],
                    "performance_metrics": {"avg_duration_ms": 0, "success_rate": 0},
                    "use_count": 0
                }
            ]
        }
        
        self._save_collaborations(default_collabs)
        return default_collabs
    
    def _save_collaborations(self, collabs: Dict):
        """保存协作配置"""
        collabs["updated_at"] = datetime.now().isoformat()
        with open(self.collaborations_file, 'w', encoding='utf-8') as f:
            json.dump(collabs, f, ensure_ascii=False, indent=2)
    
    def get_collaboration(self, collab_id: str) -> Optional[Dict]:
        """获取协作配置"""
        for collab in self.collaborations.get("collaborations", []):
            if collab.get("collaboration_id") == collab_id:
                return collab
        return None
    
    def find_applicable_collaboration(self, 
                                        problem_type: str,
                                        complexity: int) -> Optional[Dict]:
        """查找适用的协作配置"""
        applicable = []
        
        for collab in self.collaborations.get("collaborations", []):
            conditions = collab.get("conditions", [])
            
            # 检查条件匹配
            if "复杂问题" in conditions and complexity >= 7:
                applicable.append(collab)
            elif "简单问题" in conditions and complexity <= 3:
                applicable.append(collab)
            elif not conditions:
                applicable.append(collab)
        
        # 按使用次数排序
        applicable.sort(key=lambda x: x.get("use_count", 0), reverse=True)
        
        return applicable[0] if applicable else None
    
    def execute_collaboration(self, 
                               collab_id: str,
                               executors: Dict[str, Any]) -> Dict:
        """
        执行协作流程
        
        Args:
            collab_id: 协作ID
            executors: 模型执行器字典 {model_name: executor_func}
            
        Returns:
            执行结果
        """
        collab = self.get_collaboration(collab_id)
        if not collab:
            return {"status": "error", "message": "协作配置不存在"}
        
        results = []
        workflow = sorted(collab.get("workflow", []), key=lambda x: x.get("order", 0))
        
        for step in workflow:
            model = step.get("model")
            action = step.get("action")
            
            if model in executors:
                try:
                    result = executors[model](action)
                    results.append({
                        "step": step.get("order"),
                        "model": model,
                        "action": action,
                        "result": result,
                        "success": True
                    })
                except Exception as e:
                    results.append({
                        "step": step.get("order"),
                        "model": model,
                        "action": action,
                        "error": str(e),
                        "success": False
                    })
            else:
                results.append({
                    "step": step.get("order"),
                    "model": model,
                    "action": action,
                    "error": f"执行器不存在: {model}",
                    "success": False
                })
        
        # 更新使用计数
        collab["use_count"] = collab.get("use_count", 0) + 1
        self._save_collaborations(self.collaborations)
        
        return {
            "status": "completed",
            "collaboration": collab.get("name"),
            "results": results,
            "overall_success": all(r.get("success") for r in results)
        }
    
    def create_collaboration(self, collab: Dict) -> Optional[str]:
        """创建新协作配置"""
        try:
            collab["collaboration_id"] = collab.get("collaboration_id") or f"collab_{datetime.now().strftime('%Y%m%d%H%M%S')}"
            collab["use_count"] = 0
            collab["performance_metrics"] = {}
            
            self.collaborations["collaborations"].append(collab)
            self._save_collaborations(self.collaborations)
            
            return collab["collaboration_id"]
        except Exception:
            return None


class AdvancedAnalyzer:
    """高级分析器"""
    
    def __init__(self):
        """初始化高级分析器"""
        pass
    
    def analyze_success_factors(self, sessions: List[Dict]) -> Dict:
        """
        分析成功因素
        
        Args:
            sessions: 会话列表
            
        Returns:
            成功因素分析
        """
        successful = [s for s in sessions if s.get("success", False)]
        failed = [s for s in sessions if not s.get("success", False)]
        
        # 分析模型成功率
        model_success = defaultdict(lambda: {"success": 0, "total": 0})
        for s in sessions:
            model = s.get("selected_model", "unknown")
            model_success[model]["total"] += 1
            if s.get("success", False):
                model_success[model]["success"] += 1
        
        model_rates = {
            model: data["success"] / data["total"] 
            for model, data in model_success.items() 
            if data["total"] > 0
        }
        
        return {
            "total_sessions": len(sessions),
            "successful_sessions": len(successful),
            "failed_sessions": len(failed),
            "overall_success_rate": len(successful) / len(sessions) if sessions else 0,
            "model_success_rates": model_rates,
            "best_model": max(model_rates, key=model_rates.get) if model_rates else None,
            "success_factors": self._extract_success_factors(successful),
            "failure_factors": self._extract_failure_factors(failed)
        }
    
    def _extract_success_factors(self, successful: List[Dict]) -> List[str]:
        """提取成功因素"""
        factors = []
        
        # 检查高评分会话
        high_rated = [s for s in successful if s.get("user_rating", 0) >= 4]
        if high_rated:
            factors.append(f"用户评分高的会话 ({len(high_rated)}个) 通常成功")
        
        # 检查平均耗时
        if successful:
            avg_duration = statistics.mean([s.get("duration_ms", 0) for s in successful])
            factors.append(f"成功会话平均耗时: {avg_duration:.0f}ms")
        
        return factors
    
    def _extract_failure_factors(self, failed: List[Dict]) -> List[str]:
        """提取失败因素"""
        factors = []
        
        if failed:
            failed_models = defaultdict(int)
            for s in failed:
                model = s.get("selected_model", "unknown")
                failed_models[model] += 1
            
            if failed_models:
                worst_model = max(failed_models, key=failed_models.get)
                factors.append(f"{worst_model} 模型失败次数最多 ({failed_models[worst_model]}次)")
        
        return factors
    
    def generate_optimization_suggestions(self, analysis: Dict) -> List[str]:
        """生成优化建议"""
        suggestions = []
        
        # 基于成功率建议
        model_rates = analysis.get("model_success_rates", {})
        if model_rates:
            worst_model = min(model_rates, key=model_rates.get)
            worst_rate = model_rates[worst_model]
            
            if worst_rate < 0.5:
                suggestions.append(f"考虑优化 {worst_model} 模型，成功率仅为 {worst_rate*100:.0f}%")
                suggestions.append(f"建议检查该模型的适用条件和阶段配置")
        
        # 基于总体成功率
        overall = analysis.get("overall_success_rate", 0)
        if overall < 0.7:
            suggestions.append("总体成功率较低，建议增加记忆库中的成功案例")
            suggestions.append("考虑使用更简单的问题分解方法")
        elif overall > 0.9:
            suggestions.append("成功率很高！可以考虑尝试更复杂的问题")
        
        # 基于会话量
        total = analysis.get("total_sessions", 0)
        if total < 10:
            suggestions.append("样本量较小，建议继续积累更多会话数据")
        elif total > 100:
            suggestions.append("已有足够的样本数据，可以进行深度模式分析")
        
        return suggestions
    
    def predict_outcome(self, 
                        problem: str, 
                        model: str,
                        historical_data: List[Dict]) -> Dict:
        """
        预测结果
        
        Args:
            problem: 问题描述
            model: 选择的模型
            historical_data: 历史数据
            
        Returns:
            预测结果
        """
        # 查找相似历史案例
        similar = []
        for session in historical_data:
            if model == session.get("selected_model"):
                similar.append(session)
        
        if not similar:
            return {
                "predicted_success": 0.5,
                "confidence": "low",
                "reason": "缺乏历史数据",
                "suggestion": "建议从简单问题开始积累数据"
            }
        
        # 计算成功率
        successful = [s for s in similar if s.get("success", False)]
        success_rate = len(successful) / len(similar) if similar else 0.5
        
        # 评估置信度
        confidence = "high" if len(similar) >= 20 else "medium" if len(similar) >= 5 else "low"
        
        # 生成建议
        suggestions = []
        if success_rate < 0.5:
            suggestions.append("建议更换其他模型")
        elif success_rate > 0.8:
            suggestions.append("该模型对此类问题表现良好")
        
        return {
            "predicted_success": success_rate,
            "confidence": confidence,
            "similar_cases": len(similar),
            "historical_success_rate": success_rate,
            "suggestions": suggestions
        }


def get_continuous_learning() -> ContinuousLearning:
    """获取持续学习系统实例"""
    return ContinuousLearning()


def get_model_collaborator() -> ModelCollaborator:
    """获取多模型协同管理器实例"""
    return ModelCollaborator()


def get_advanced_analyzer() -> AdvancedAnalyzer:
    """获取高级分析器实例"""
    return AdvancedAnalyzer()
