"""
Task Decomposition Agent for Task Manager MCP Servers

Adapted from the original project's TaskDecompositionAgent for use in MCP servers.
This agent evaluates task complexity and decomposes complex tasks into manageable subtasks.
"""

import requests
import json
from typing import Dict, List, Optional
import re
import os


class TaskDecompositionAgent:
    """
    一个专门的Agent，负责评估单个任务的复杂性，并在必要时将其分解为更小的子任务。
    """
    
    SYSTEM_PROMPT_TEMPLATE = """
You are a task decomposition expert. Your job is to break a single complex task into 3 to 5 high-quality, meaningful sub-tasks.

**CRITICAL RULES:**
1.  **Analyze Complexity:** First, decide if the input task is complex enough to need decomposition. A task like "Implement user authentication" is complex. A task like "Create a login button component" is NOT.
2.  **Decomposition Rules:**
    -   If the task is **NOT complex**, you MUST return an empty list: `{ "sub_tasks": [] }`.
    -   If the task **IS complex**, you MUST break it into exactly 3 to 5 sub-tasks.
    -   Each sub-task should be a significant piece of work, not a tiny step. A sub-task's description can list smaller steps within it.
3.  **What to AVOID:**
    -   **NEVER** create sub-tasks for trivial steps like "create a file", "add a dependency", or "write a single function". Group these into a larger, logical task.
4.  **Output Format:** Your response MUST be a single JSON object with one key, `"sub_tasks"`, containing the list of sub-task objects.

---
**EXAMPLE 1: Complex Task (Correct Decomposition)**

*Input:*
`{ "name": "Build a complete user profile page" }`

*Output:*
```json
{
  "sub_tasks": [
    {
      "id": "placeholder-1",
      "name": "Design Profile Page UI/UX",
      "description": "Create wireframes and mockups for viewing and editing profiles."
    },
    {
      "id": "placeholder-2",
      "name": "Develop Backend API for Profiles",
      "description": "Build GET and PUT endpoints for profile data. Includes data validation."
    },
    {
      "id": "placeholder-3",
      "name": "Implement Frontend Profile Page UI",
      "description": "Develop the React components for the page, including the form for editing."
    }
  ]
}
```
---
**EXAMPLE 2: Simple Task (Correctly Not Decomposed)**

*Input:*
`{ "name": "Add a 'Sort by Price' button to the products page" }`

*Output:*
```json
{
  "sub_tasks": []
}
```
"""

    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None):
        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 environment variable (OPENAI_API_KEY).")

    def decompose_task(self, parent_task: Dict, existing_task_ids: List[str]) -> List[Dict]:
        """
        Takes a parent task and returns a list of sub-tasks if decomposition is needed.
        """
        print(f"\n--- 任务拆解Agent (父任务: {parent_task.get('id', 'unknown')}) ---")
        print(f"⏳ 正在评估任务 '{parent_task.get('name', 'unnamed')}' 的复杂度并决定是否拆解...")

        base_id_num = len(existing_task_ids) + 1

        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}"
        }

        # Ensure proper encoding for the payload
        try:
            encoded_system_prompt = self.SYSTEM_PROMPT_TEMPLATE.encode('utf-8', errors='replace').decode('utf-8')
            encoded_task_content = json.dumps(parent_task, indent=2, ensure_ascii=False).encode('utf-8', errors='replace').decode('utf-8')
        except (UnicodeEncodeError, UnicodeDecodeError):
            encoded_system_prompt = self.SYSTEM_PROMPT_TEMPLATE
            encoded_task_content = json.dumps(parent_task, indent=2, ensure_ascii=False)

        payload = {
            "model": "gpt-4o",
            "messages": [
                {"role": "system", "content": encoded_system_prompt},
                {"role": "user", "content": encoded_task_content}
            ],
            "temperature": 0.1,
        }

        try:
            # Use json parameter instead of data for proper encoding
            response = requests.post(
                api_url,
                headers=headers,
                json=payload,  # This automatically handles encoding
                timeout=120
            )
            response.raise_for_status()

            response_data = response.json()
            content = response_data.get("choices", [{}])[0].get("message", {}).get("content", "")

            # Ensure proper encoding for the response
            try:
                content = content.encode('utf-8', errors='replace').decode('utf-8')
            except (UnicodeEncodeError, UnicodeDecodeError):
                pass
            
            # Extract JSON from response
            json_match = re.search(r'\{.*\}', content, re.DOTALL)
            if not json_match:
                print("❌ 无法从LLM响应中提取JSON。")
                return []
            
            try:
                result = json.loads(json_match.group())
                sub_tasks = result.get("sub_tasks", [])
                
                if not sub_tasks:
                    print("✅ 任务被判定为不需要拆解。")
                    return []
                
                # Assign proper IDs to sub-tasks
                for i, sub_task in enumerate(sub_tasks):
                    if sub_task.get("id", "").startswith("placeholder"):
                        sub_task["id"] = f"T-{base_id_num + i:03d}"
                
                print(f"✅ 任务已拆解为 {len(sub_tasks)} 个子任务。")
                return sub_tasks
                
            except json.JSONDecodeError as e:
                error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')
                print(f"❌ JSON解析失败: {error_msg}")
                return []
                
        except requests.exceptions.RequestException as e:
            error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')
            print(f"❌ API调用失败: {error_msg}")
            return []
        except Exception as e:
            error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')
            print(f"❌ 未知错误: {error_msg}")
            return []