"""
Dependency Manager for Task Manager MCP Servers

Adapted from the original project's DependencyManager for use in MCP servers.
This agent analyzes task dependencies and manages task relationships.
"""

import os
import json
import requests
import re
from typing import Dict, List, Set, Optional, Any
from datetime import datetime


class DependencyManager:
    """
    一个Agent，负责分析任务之间的依赖关系，并根据任务完成情况更新状态。
    """
    
    def __init__(self, tasks: Dict[str, Any], api_key: Optional[str] = None, base_url: Optional[str] = None):
        """
        初始化依赖管理器。

        Args:
            tasks: 任务字典
            api_key: The OpenAI API key.
            base_url: The base URL for the OpenAI API.
        """
        self.tasks = tasks
        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
        self.base_url = base_url or os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
        if not self.api_key:
            raise ValueError("API key must be provided or set as an environment variable.")

    def _call_llm_for_json(self, payload: Dict) -> Optional[Dict]:
        """
        一个通用的、使用requests库调用LLM并期望返回JSON的辅助函数。
        """
        api_url = f"{self.base_url}/chat/completions"
        if not self.base_url.endswith('/v1'):
            api_url = f"{self.base_url}/v1/chat/completions"
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        try:
            response = requests.post(api_url, headers=headers, json=payload, timeout=120)
            response.raise_for_status()

            response_data = response.json()
            response_content = response_data["choices"][0]["message"]["content"]
            
            if not response_content:
                print("❌ LLM返回了空内容。")
                return None
            
            match = re.search(r"\{[\s\S]*\}", response_content)
            if not match:
                print(f"❌ 在LLM响应中找不到JSON对象: {response_content}")
                return None
            
            return json.loads(match.group(0))

        except Exception as e:
            print(f"❌ 调用LLM时发生错误: {e}")
            return None

    def add_dependencies_with_ai(self, tasks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        使用AI分析任务列表，并添加依赖关系。
        """
        print("\n--- 依赖分析Agent ---")
        print("⏳ 正在调用LLM，分析任务间的依赖关系...")

        prompt = self._create_dependency_prompt(tasks)
        
        payload = {
            "model": "gpt-4o",
            "messages": [
                {"role": "system", "content": prompt},
                {"role": "user", "content": json.dumps(tasks, indent=2, ensure_ascii=False)}
            ],
            "temperature": 0.2,
        }

        tasks_with_deps_data = self._call_llm_for_json(payload)

        if not tasks_with_deps_data:
            print(f"❌ 解析失败或LLM调用失败。")
            return tasks

        if isinstance(tasks_with_deps_data, dict):
            for key, value in tasks_with_deps_data.items():
                if isinstance(value, list) and len(value) == len(tasks):
                    print(f"✅ 成功从键 '{key}' 中提取带依赖的任务列表。")
                    return value
        
        print(f"❌ 解析失败：LLM返回的JSON中未找到格式正确的任务列表。")
        return tasks
    
    def _create_dependency_prompt(self, tasks: List[Dict]) -> str:
        """为依赖分析创建一个专用的Prompt，现在也包括优先级和成本估算。"""
        return f"""
你是一个专业的项目管理AI，负责分析任务之间的依赖关系。你的任务是：

1. **分析依赖关系**：识别哪些任务必须在其他任务完成后才能开始
2. **设置优先级**：基于依赖关系、重要性和紧急程度设置任务优先级（1-5，1最高）
3. **估算成本**：评估每个任务的时间成本、重要性、收益和边际成本

**依赖关系规则**：
- 如果任务A的输出是任务B的输入，则B依赖于A
- 如果任务A必须在任务B之前完成，则B依赖于A
- 避免循环依赖
- 只添加真正必要的依赖关系

**优先级设置**（1-5）：
- 1（最高）：关键路径上的任务，阻塞其他任务
- 2（高）：重要且有依赖关系的任务
- 3（中）：一般重要性的任务
- 4（低）：可以延后的任务
- 5（最低）：非关键任务

**成本估算**：
- time_cost: 预估完成时间（小时）
- importance: 任务重要性（1-5，5最重要）
- benefit: 预期收益（1-10，10最高）
- marginal_cost: 边际成本（1-5，5最高）

**输出格式**：返回一个JSON对象，包含一个"tasks"键，值为更新后的任务列表。每个任务必须包含：
- 原有的所有字段
- dependencies: 依赖的任务ID列表
- priority: 优先级（1-5）
- time_cost: 时间成本（小时）
- importance: 重要性（1-5）
- benefit: 收益（1-10）
- marginal_cost: 边际成本（1-5）

示例输出：
```json
{{
  "tasks": [
    {{
      "id": "task1",
      "name": "设计数据库架构",
      "description": "...",
      "dependencies": [],
      "priority": 1,
      "time_cost": 8.0,
      "importance": 5,
      "benefit": 9,
      "marginal_cost": 2
    }},
    {{
      "id": "task2", 
      "name": "实现用户认证",
      "description": "...",
      "dependencies": ["task1"],
      "priority": 2,
      "time_cost": 12.0,
      "importance": 4,
      "benefit": 8,
      "marginal_cost": 3
    }}
  ]
}}
```

请分析以下任务列表：
"""

    def get_ready_tasks(self) -> List[str]:
        """
        获取所有可以开始执行的任务（即所有依赖都已完成的任务）。
        """
        ready_tasks = []
        
        for task_id, task in self.tasks.items():
            if task.get('status') in ['todo', 'pending']:
                dependencies = task.get('dependencies', [])
                
                # 检查所有依赖是否都已完成
                all_deps_completed = True
                for dep_id in dependencies:
                    dep_task = self.tasks.get(dep_id)
                    if not dep_task or dep_task.get('status') != 'completed':
                        all_deps_completed = False
                        break
                
                if all_deps_completed:
                    ready_tasks.append(task_id)
        
        return ready_tasks

    def get_blocked_tasks(self) -> List[str]:
        """
        获取所有被阻塞的任务（即有未完成依赖的任务）。
        """
        blocked_tasks = []
        
        for task_id, task in self.tasks.items():
            if task.get('status') in ['todo', 'pending']:
                dependencies = task.get('dependencies', [])
                
                # 检查是否有未完成的依赖
                has_incomplete_deps = False
                for dep_id in dependencies:
                    dep_task = self.tasks.get(dep_id)
                    if not dep_task or dep_task.get('status') != 'completed':
                        has_incomplete_deps = True
                        break
                
                if has_incomplete_deps:
                    blocked_tasks.append(task_id)
        
        return blocked_tasks

    def validate_dependencies(self) -> Dict[str, List[str]]:
        """
        验证依赖关系，检查循环依赖和无效依赖。
        
        Returns:
            Dict with 'circular' and 'invalid' keys containing lists of problematic task IDs
        """
        issues = {'circular': [], 'invalid': []}
        
        # 检查无效依赖（依赖不存在的任务）
        for task_id, task in self.tasks.items():
            dependencies = task.get('dependencies', [])
            for dep_id in dependencies:
                if dep_id not in self.tasks:
                    issues['invalid'].append(f"{task_id} depends on non-existent task {dep_id}")
        
        # 检查循环依赖（简化版本）
        def has_circular_dependency(task_id: str, visited: Set[str], path: Set[str]) -> bool:
            if task_id in path:
                return True
            if task_id in visited:
                return False
            
            visited.add(task_id)
            path.add(task_id)
            
            task = self.tasks.get(task_id, {})
            dependencies = task.get('dependencies', [])
            
            for dep_id in dependencies:
                if has_circular_dependency(dep_id, visited, path):
                    return True
            
            path.remove(task_id)
            return False
        
        visited = set()
        for task_id in self.tasks:
            if task_id not in visited:
                if has_circular_dependency(task_id, visited, set()):
                    issues['circular'].append(task_id)
        
        return issues
