#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
思维模型增强器 - 高级功能模块
Thinking Model Enhancer - Advanced Features Module

实现高级功能：性能追踪、思维可视化、批量处理、自定义配置
"""

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


@dataclass
class PerformanceMetrics:
    """性能指标"""
    total_runs: int = 0
    successful_runs: int = 0
    failed_runs: int = 0
    total_duration_ms: float = 0
    avg_duration_ms: float = 0
    min_duration_ms: float = float('inf')
    max_duration_ms: float = 0
    by_model: Dict[str, Dict] = field(default_factory=dict)
    by_problem_type: Dict[str, Dict] = field(default_factory=dict)
    recent_runs: List[Dict] = field(default_factory=list)
    
    def to_dict(self) -> Dict:
        return {
            "total_runs": self.total_runs,
            "successful_runs": self.successful_runs,
            "failed_runs": self.failed_runs,
            "success_rate": self.successful_runs / self.total_runs if self.total_runs > 0 else 0,
            "total_duration_ms": self.total_duration_ms,
            "avg_duration_ms": self.avg_duration_ms,
            "min_duration_ms": self.min_duration_ms if self.min_duration_ms != float('inf') else 0,
            "max_duration_ms": self.max_duration_ms,
            "by_model": self.by_model,
            "by_problem_type": self.by_problem_type,
            "recent_runs": self.recent_runs[-50:]  # 最近50条
        }


class PerformanceTracker:
    """性能追踪器"""
    
    def __init__(self, storage_dir: Optional[str] = None):
        """初始化性能追踪器"""
        if storage_dir is None:
            self.storage_dir = Path.home() / ".claude" / "thinking_models" / "performance"
        else:
            self.storage_dir = Path(storage_dir)
        
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self.metrics_file = self.storage_dir / "metrics.json"
        
        # 加载现有指标
        self.metrics = self._load_metrics()
        
        # 实时追踪
        self.current_run_start: Optional[float] = None
        self.current_run_data: Dict = {}
    
    def _load_metrics(self) -> PerformanceMetrics:
        """加载性能指标"""
        if self.metrics_file.exists():
            with open(self.metrics_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
                metrics = PerformanceMetrics()
                metrics.total_runs = data.get('total_runs', 0)
                metrics.successful_runs = data.get('successful_runs', 0)
                metrics.failed_runs = data.get('failed_runs', 0)
                metrics.total_duration_ms = data.get('total_duration_ms', 0)
                metrics.avg_duration_ms = data.get('avg_duration_ms', 0)
                metrics.min_duration_ms = data.get('min_duration_ms', 0)
                metrics.max_duration_ms = data.get('max_duration_ms', 0)
                metrics.by_model = data.get('by_model', {})
                metrics.by_problem_type = data.get('by_problem_type', {})
                metrics.recent_runs = data.get('recent_runs', [])
                return metrics
        return PerformanceMetrics()
    
    def _save_metrics(self):
        """保存性能指标"""
        with open(self.metrics_file, 'w', encoding='utf-8') as f:
            json.dump(self.metrics.to_dict(), f, ensure_ascii=False, indent=2)
    
    def start_run(self, model: str, problem_type: str, problem: str):
        """开始一次运行追踪"""
        self.current_run_start = time.time()
        self.current_run_data = {
            "model": model,
            "problem_type": problem_type,
            "problem": problem[:100],
            "start_time": datetime.now().isoformat(),
            "stages": []
        }
    
    def log_stage(self, stage_name: str, duration_ms: float):
        """记录阶段信息"""
        if self.current_run_data:
            self.current_run_data["stages"].append({
                "stage": stage_name,
                "duration_ms": duration_ms
            })
    
    def end_run(self, success: bool, result_summary: str = ""):
        """结束运行追踪"""
        if self.current_run_start is None:
            return
        
        duration_ms = (time.time() - self.current_run_start) * 1000
        
        # 更新总体指标
        self.metrics.total_runs += 1
        if success:
            self.metrics.successful_runs += 1
        else:
            self.metrics.failed_runs += 1
        
        self.metrics.total_duration_ms += duration_ms
        self.metrics.avg_duration_ms = self.metrics.total_duration_ms / self.metrics.total_runs
        self.metrics.min_duration_ms = min(self.metrics.min_duration_ms, duration_ms)
        self.metrics.max_duration_ms = max(self.metrics.max_duration_ms, duration_ms)
        
        # 更新模型指标
        model = self.current_run_data.get("model", "unknown")
        if model not in self.metrics.by_model:
            self.metrics.by_model[model] = {
                "runs": 0, "success": 0, "total_duration": 0
            }
        self.metrics.by_model[model]["runs"] += 1
        if success:
            self.metrics.by_model[model]["success"] += 1
        self.metrics.by_model[model]["total_duration"] += duration_ms
        
        # 更新问题类型指标
        problem_type = self.current_run_data.get("problem_type", "unknown")
        if problem_type not in self.metrics.by_problem_type:
            self.metrics.by_problem_type[problem_type] = {
                "runs": 0, "success": 0
            }
        self.metrics.by_problem_type[problem_type]["runs"] += 1
        if success:
            self.metrics.by_problem_type[problem_type]["success"] += 1
        
        # 添加到最近记录
        self.metrics.recent_runs.append({
            "timestamp": datetime.now().isoformat(),
            "model": model,
            "problem_type": problem_type,
            "duration_ms": duration_ms,
            "success": success,
            "stages": len(self.current_run_data.get("stages", [])),
            "summary": result_summary[:50]
        })
        
        # 保存并重置
        self._save_metrics()
        self.current_run_start = None
        self.current_run_data = {}
    
    def get_metrics(self) -> PerformanceMetrics:
        """获取性能指标"""
        return self.metrics
    
    def get_model_performance(self, model: str) -> Dict:
        """获取特定模型的性能"""
        return self.metrics.by_model.get(model, {
            "runs": 0, "success": 0, "total_duration": 0
        })
    
    def get_summary_report(self) -> str:
        """生成性能摘要报告"""
        m = self.metrics
        success_rate = m.successful_runs / m.total_runs if m.total_runs > 0 else 0
        
        lines = [
            "📊 思维模型性能报告",
            "=" * 50,
            "",
            f"📈 总体统计:",
            f"   总运行次数: {m.total_runs}",
            f"   成功: {m.successful_runs} | 失败: {m.failed_runs}",
            f"   成功率: {success_rate*100:.1f}%",
            f"   平均耗时: {m.avg_duration_ms:.1f}ms",
            f"   最短耗时: {m.min_duration_ms:.1f}ms" if m.min_duration_ms != float('inf') else "",
            f"   最长耗时: {m.max_duration_ms:.1f}ms",
        ]
        
        if m.by_model:
            lines.extend(["", f"📂 按模型:"])
            for model, data in sorted(m.by_model.items(), key=lambda x: x[1]["runs"], reverse=True):
                model_success = data["success"] / data["runs"] if data["runs"] > 0 else 0
                lines.append(f"   • {model}: {data['runs']}次 (成功率{ model_success*100:.0f}%)")
        
        if m.by_problem_type:
            lines.extend(["", f"🏷️ 按问题类型:"])
            for ptype, data in sorted(m.by_problem_type.items(), key=lambda x: x[1]["runs"], reverse=True):
                lines.append(f"   • {ptype}: {data['runs']}次")
        
        return "\n".join(filter(None, lines))


class ThinkingVisualizer:
    """思维可视化器"""
    
    def __init__(self):
        """初始化可视化器"""
        pass
    
    def format_thinking_process(self, 
                                 problem: str,
                                 analysis: Dict,
                                 model: str,
                                 stages: List[Dict],
                                 result: Dict) -> str:
        """
        格式化思维过程为可读报告
        
        Args:
            problem: 问题描述
            analysis: 问题分析结果
            model: 使用的模型
            stages: 各阶段处理结果
            result: 最终结果
            
        Returns:
            格式化的报告字符串
        """
        lines = [
            f"🧠 思维过程报告",
            "=" * 60,
            "",
            f"📝 问题: {problem[:100]}{'...' if len(problem) > 100 else ''}",
            "",
            f"🔍 问题分析:",
            f"   类型: {analysis.get('type', '未知')}",
            f"   复杂度: {analysis.get('complexity', 'N/A')}/10",
            f"   置信度: {analysis.get('confidence', 'N/A')}",
            "",
            f"🎯 使用模型: {model}",
            "",
            f"🔄 处理过程 ({len(stages)} 阶段):"
        ]
        
        for i, stage in enumerate(stages, 1):
            lines.append(f"   {i}. {stage.get('name', 'Unknown')}")
            if stage.get("findings"):
                for finding in stage["findings"][:3]:
                    lines.append(f"      • {finding}")
            if stage.get("duration_ms"):
                lines.append(f"      ⏱️ {stage['duration_ms']:.1f}ms")
        
        lines.extend([
            "",
            f"✅ 处理结果:",
            f"   置信度: {result.get('confidence', 'N/A')}",
            f"   耗时: {result.get('duration_ms', 0):.1f}ms",
        ])
        
        if result.get("recommendations"):
            lines.extend([
                "",
                f"💡 建议:",
                *[f"   • {rec}" for rec in result["recommendations"][:5]]
            ])
        
        return "\n".join(lines)
    
    def generate_text_chart(self, 
                            data: List[Dict], 
                            value_key: str = "success_rate",
                            label_key: str = "label") -> str:
        """
        生成简单的文本图表
        
        Args:
            data: 数据列表
            value_key: 数值字段名
            label_key: 标签字段名
            
        Returns:
            文本图表字符串
        """
        if not data:
            return "暂无数据 📊"
        
        max_value = max(d.get(value_key, 0) for d in data)
        max_label_len = max(len(d.get(label_key, "")) for d in data)
        
        lines = []
        for item in data:
            label = item.get(label_key, "")[:20]
            value = item.get(value_key, 0)
            
            # 计算进度条长度
            bar_length = int((value / max_value) * 20) if max_value > 0 else 0
            bar = "█" * bar_length + "░" * (20 - bar_length)
            
            # 格式化行
            if isinstance(value, float):
                if value > 1:
                    line = f"{label:<{max_label_len}} {bar} {value:.0f}"
                else:
                    line = f"{label:<{max_label_len}} {bar} {value*100:.1f}%"
            else:
                line = f"{label:<{max_label_len}} {bar} {value}"
            
            lines.append(line)
        
        return "\n".join(lines)


class BatchProcessor:
    """批量处理器"""
    
    def __init__(self, 
                 processor_func: Callable,
                 max_concurrent: int = 3,
                 timeout_seconds: int = 60):
        """
        初始化批量处理器
        
        Args:
            processor_func: 处理函数
            max_concurrent: 最大并发数
            timeout_seconds: 超时时间
        """
        self.processor = processor_func
        self.max_concurrent = max_concurrent
        self.timeout = timeout_seconds
        
        self.results: List[Dict] = []
        self.stats = {
            "total": 0,
            "completed": 0,
            "failed": 0,
            "start_time": None,
            "end_time": None
        }
    
    def process(self, items: List[Dict]) -> List[Dict]:
        """
        批量处理项目
        
        Args:
            items: 项目列表，每个项目包含id和data
            
        Returns:
            处理结果列表
        """
        self.stats["total"] = len(items)
        self.stats["start_time"] = datetime.now().isoformat()
        
        results = []
        for item in items:
            try:
                result = self.processor(item["data"])
                results.append({
                    "id": item.get("id", ""),
                    "status": "success",
                    "result": result
                })
                self.stats["completed"] += 1
            except Exception as e:
                results.append({
                    "id": item.get("id", ""),
                    "status": "failed",
                    "error": str(e)
                })
                self.stats["failed"] += 1
        
        self.stats["end_time"] = datetime.now().isoformat()
        self.results = results
        
        return results
    
    def get_stats(self) -> Dict:
        """获取处理统计"""
        duration = None
        if self.stats["start_time"] and self.stats["end_time"]:
            start = datetime.fromisoformat(self.stats["start_time"])
            end = datetime.fromisoformat(self.stats["end_time"])
            duration = (end - start).total_seconds()
        
        return {
            **self.stats,
            "duration_seconds": duration,
            "success_rate": self.stats["completed"] / self.stats["total"] if self.stats["total"] > 0 else 0
        }


class ModelConfigManager:
    """模型配置管理器"""
    
    def __init__(self, config_dir: Optional[str] = None):
        """初始化配置管理器"""
        if config_dir is None:
            self.config_dir = Path.home() / ".claude" / "thinking_models" / "config"
        else:
            self.config_dir = Path(config_dir)
        
        self.config_dir.mkdir(parents=True, exist_ok=True)
        self.config_file = self.config_dir / "models.json"
        
        # 加载或创建默认配置
        self.config = self._load_or_create_config()
    
    def _load_or_create_config(self) -> Dict:
        """加载或创建默认配置"""
        if self.config_file.exists():
            with open(self.config_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        
        # 默认配置
        default_config = {
            "version": "1.0",
            "updated_at": datetime.now().isoformat(),
            "models": {
                "research_mode": {
                    "enabled": True,
                    "timeout_seconds": 120,
                    "max_retries": 3,
                    "priority": 1
                },
                "diagnostic_mode": {
                    "enabled": True,
                    "timeout_seconds": 90,
                    "max_retries": 2,
                    "priority": 2
                },
                "generic_pipeline": {
                    "enabled": True,
                    "timeout_seconds": 60,
                    "max_retries": 3,
                    "priority": 10
                }
            },
            "general": {
                "auto_store_results": True,
                "confidence_threshold": 0.6,
                "enable_visualization": True,
                "max_history_items": 100
            }
        }
        
        self._save_config(default_config)
        return default_config
    
    def _save_config(self, config: Dict):
        """保存配置"""
        config["updated_at"] = datetime.now().isoformat()
        with open(self.config_file, 'w', encoding='utf-8') as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
    
    def get_model_config(self, model_name: str) -> Optional[Dict]:
        """获取模型配置"""
        return self.config.get("models", {}).get(model_name)
    
    def set_model_config(self, model_name: str, settings: Dict):
        """设置模型配置"""
        if model_name not in self.config["models"]:
            self.config["models"][model_name] = {}
        
        self.config["models"][model_name].update(settings)
        self._save_config(self.config)
    
    def get_general_config(self) -> Dict:
        """获取通用配置"""
        return self.config.get("general", {})
    
    def set_general_config(self, settings: Dict):
        """设置通用配置"""
        self.config["general"].update(settings)
        self._save_config(self.config)
    
    def is_model_enabled(self, model_name: str) -> bool:
        """检查模型是否启用"""
        model_config = self.get_model_config(model_name)
        return model_config.get("enabled", True) if model_config else True


def get_performance_tracker() -> PerformanceTracker:
    """获取性能追踪器实例"""
    return PerformanceTracker()


def get_visualizer() -> ThinkingVisualizer:
    """获取可视化器实例"""
    return ThinkingVisualizer()


def get_batch_processor(processor_func: Callable) -> BatchProcessor:
    """获取批量处理器实例"""
    return BatchProcessor(processor_func)


def get_config_manager() -> ModelConfigManager:
    """获取配置管理器实例"""
    return ModelConfigManager()
