mnemo_cards/tools/agent/planning_agent.py

"""
Planning Agent - Analyzes project and generates task lists.
"""
import sys
import json
from pathlib import Path
from datetime import datetime, timezone
from typing import Dict, Any, List

from config import AgentConfig
from cursor_cli_wrapper import CursorCLI, CursorResultStatus
from task_manager import Task


class PlanningAgent:
    """Planning agent that generates task lists."""
    
    def __init__(self, config: AgentConfig):
        self.config = config
        self.cursor_cli = CursorCLI(
            project_root=config.project_root,
            api_key=config.cursor_api_key,
            model=config.cursor_model,
            verbose=True
        )
    
    def run(self) -> int:
        """
        Generate task list for component.
        Returns: 0 on success, 1 on error
        """
        print(f"🎯 Planning Agent for component: {self.config.component}")
        print(f"📁 Project root: {self.config.project_root}")
        
        try:
            # Build planning prompt
            prompt = self._build_planning_prompt()
            
            # Run cursor agent to analyze and plan
            print("\n🤖 Running Cursor Planning Agent...")
            result = self.cursor_cli.run_agent(
                task_description=prompt,
                force=True,  # Allow writing task_list.json
                max_iterations=3
            )
            
            if result.status != CursorResultStatus.SUCCESS:
                print(f"❌ Planning agent failed: {result.error}")
                return 1
            
            # Verify task list was generated
            if not self.config.task_list_path.exists():
                print(f"❌ Task list not generated: {self.config.task_list_path}")
                return 1
            
            # Validate and summarize task list
            task_list = self._load_and_validate_task_list()
            
            if not task_list:
                print("❌ Invalid task list generated")
                return 1
            
            # Print summary
            self._print_summary(task_list)
            
            # Create summary commit
            self._commit_task_list()
            
            print("\n✅ Planning completed successfully")
            return 0
        
        except Exception as e:
            print(f"\n❌ Planning error: {e}")
            import traceback
            traceback.print_exc()
            return 1
    
    def _build_planning_prompt(self) -> str:
        """Build comprehensive planning prompt."""
        # Load planning prompt template
        prompt_file = self.config.prompts_dir / "planning_prompt.md"
        
        if prompt_file.exists():
            with open(prompt_file, 'r') as f:
                template = f.read()
        else:
            template = "You are an AI planning agent. Analyze the project and create a task list.\n\n"
        
        # Add context from existing files
        context = self._gather_context()
        
        prompt = f"""{template}

## PLANNING REQUEST

Generate a task list for component: **{self.config.component}**

### Current Context:

{context}

### Task List Requirements:

1. Analyze the current state from tasks.md and workflow_state.md
2. Review what's already completed in agent_state.json
3. Create new tasks or update existing ones
4. Ensure proper dependencies and priorities
5. Write the complete task list to: `{self.config.task_list_path}`

### Output Format:

Write a JSON file at `{self.config.task_list_path}` with this structure:

```json
{{
  "project": "mnemo_cards_{self.config.component}",
  "component": "{self.config.component}",
  "version": "1.0",
  "generated_at": "{datetime.now(timezone.utc).isoformat()}",
  "generated_by": "planning_agent",
  "tasks": [
    {{
      "id": "TASK-001",
      "title": "Clear title",
      "priority": "high|medium|low",
      "status": "pending",
      "estimated_hours": 4.0,
      "description": "Detailed description",
      "acceptance_criteria": ["Criterion 1", "Criterion 2"],
      "dependencies": [],
      "files_to_modify": ["path/to/file.dart"],
      "component": "{self.config.component}"
    }}
  ]
}}
```

### Guidelines:

- Create 5-10 actionable tasks
- Prioritize based on business value and dependencies
- Each task should be 2-8 hours of work
- Include clear acceptance criteria
- Specify files to modify
- Set proper dependencies

Start planning now!
"""
        
        return prompt
    
    def _gather_context(self) -> str:
        """Gather context from existing files."""
        context = ""
        
        # Read tasks.md - check multiple locations
        # Priority: component-specific > global agent tasks.md
        tasks_md = self.config.component_root / "tasks.md"
        if not tasks_md.exists():
            # Fallback to global agent tasks.md
            tasks_md = self.config.project_root / "ai_docs" / "agent" / "tasks.md"
        
        if tasks_md.exists():
            with open(tasks_md, 'r') as f:
                content = f.read()
                # Truncate if too long
                if len(content) > 10000:
                    content = content[:10000] + "\n... (truncated)"
                context += f"### tasks.md:\n```\n{content}\n```\n\n"
        
        # Read workflow_state.md - component-specific only
        workflow_state = self.config.component_root / "workflow_state.md"
        
        if workflow_state.exists():
            with open(workflow_state, 'r') as f:
                content = f.read()
                if len(content) > 5000:
                    content = content[:5000] + "\n... (truncated)"
                context += f"### workflow_state.md:\n```\n{content}\n```\n\n"
        
        # Read current task list if exists
        if self.config.task_list_path.exists():
            with open(self.config.task_list_path, 'r') as f:
                current_tasks = json.load(f)
                context += f"### Current task_list.json:\n```json\n{json.dumps(current_tasks, indent=2)}\n```\n\n"
        
        # Read agent state
        if self.config.agent_state_path.exists():
            with open(self.config.agent_state_path, 'r') as f:
                state = json.load(f)
                context += f"### Current agent_state.json:\n```json\n{json.dumps(state, indent=2)}\n```\n\n"
        
        if not context:
            context = "No existing context files found. Create fresh task list based on project structure.\n"
        
        return context
    
    def _load_and_validate_task_list(self) -> Dict[str, Any]:
        """Load and validate generated task list."""
        try:
            with open(self.config.task_list_path, 'r') as f:
                task_list = json.load(f)
            
            # Validate structure
            if "tasks" not in task_list:
                print("❌ Task list missing 'tasks' field")
                return None
            
            tasks = task_list["tasks"]
            
            if not isinstance(tasks, list):
                print("❌ 'tasks' must be a list")
                return None
            
            # Validate each task
            required_fields = [
                "id", "title", "priority", "status", 
                "estimated_hours", "description", 
                "acceptance_criteria", "dependencies", 
                "files_to_modify", "component"
            ]
            
            for i, task in enumerate(tasks):
                for field in required_fields:
                    if field not in task:
                        print(f"❌ Task {i} missing required field: {field}")
                        return None
            
            return task_list
        
        except json.JSONDecodeError as e:
            print(f"❌ Invalid JSON in task list: {e}")
            return None
        except Exception as e:
            print(f"❌ Error loading task list: {e}")
            return None
    
    def _print_summary(self, task_list: Dict[str, Any]) -> None:
        """Print task list summary."""
        tasks = task_list.get("tasks", [])
        
        print(f"\n{'='*80}")
        print(f"📋 Task List Summary for {self.config.component}")
        print(f"{'='*80}")
        print(f"Total tasks: {len(tasks)}")
        
        # Count by priority
        high_count = sum(1 for t in tasks if t.get("priority") == "high")
        medium_count = sum(1 for t in tasks if t.get("priority") == "medium")
        low_count = sum(1 for t in tasks if t.get("priority") == "low")
        
        print(f"\nPriority breakdown:")
        print(f"  🔴 HIGH: {high_count}")
        print(f"  🟡 MEDIUM: {medium_count}")
        print(f"  🟢 LOW: {low_count}")
        
        # Total estimated hours
        total_hours = sum(t.get("estimated_hours", 0) for t in tasks)
        print(f"\nTotal estimated: {total_hours:.1f} hours")
        
        # List high priority tasks
        print(f"\n🔴 High Priority Tasks:")
        for task in tasks:
            if task.get("priority") == "high":
                print(f"  - {task['id']}: {task['title']} ({task['estimated_hours']}h)")
        
        print(f"\n{'='*80}")
    
    def _commit_task_list(self) -> None:
        """Commit the generated task list."""
        import subprocess
        
        # Add task list file
        subprocess.run(
            ["git", "add", str(self.config.task_list_path)],
            cwd=self.config.project_root
        )
        
        # Commit
        commit_message = f"""plan({self.config.component}): Update task list

Generated by: Planning Agent
Timestamp: {datetime.now(timezone.utc).isoformat()}
Component: {self.config.component}
"""
        
        result = subprocess.run(
            ["git", "commit", "-m", commit_message],
            cwd=self.config.project_root,
            capture_output=True,
            text=True
        )
        
        if result.returncode == 0:
            print("✅ Task list committed")
            
            # Push
            subprocess.run(
                ["git", "push"],
                cwd=self.config.project_root
            )
            print("✅ Pushed to remote")
        else:
            print("⚠️  No changes to commit (task list unchanged)")
    
    def break_down_task(self, task: Task, failure_reason: str = "") -> bool:
        """
        Break down a failed task into smaller subtasks.
        
        Args:
            task: The task that failed
            failure_reason: Reason why the task failed
        
        Returns:
            True if subtasks were successfully created, False otherwise
        """
        print(f"\n🔀 Breaking down task {task.id} into subtasks...")
        print(f"   Task: {task.title}")
        print(f"   Reason: {failure_reason}")
        
        try:
            # Build prompt for breaking down the task
            prompt = self._build_breakdown_prompt(task, failure_reason)
            
            # Run cursor agent to break down the task
            print("\n🤖 Running Cursor Agent to break down task...")
            result = self.cursor_cli.run_agent(
                task_description=prompt,
                force=True,
                max_iterations=3
            )
            
            if result.status != CursorResultStatus.SUCCESS:
                print(f"❌ Failed to break down task: {result.error}")
                return False
            
            # Load current task list before breakdown
            if not self.config.task_list_path.exists():
                print(f"❌ Task list not found: {self.config.task_list_path}")
                return False
            
            # Get original task list for comparison
            original_task_list = self._load_and_validate_task_list()
            if not original_task_list:
                print("❌ Invalid task list")
                return False
            
            original_tasks = original_task_list.get("tasks", [])
            original_task_ids = {t.get("id") for t in original_tasks}
            original_count = len(original_tasks)
            
            # Reload to get updated task list after agent execution
            task_list = self._load_and_validate_task_list()
            if not task_list:
                print("❌ Failed to reload task list after breakdown")
                return False
            
            updated_tasks = task_list.get("tasks", [])
            new_count = len(updated_tasks)
            
            # Find subtasks - they should have IDs containing the original task ID
            # or be new tasks added after the breakdown
            # Look for tasks with IDs like "TASK-001-SUB-1" or similar patterns
            subtasks = [
                t for t in updated_tasks 
                if (task.id in t.get("id", "") and t.get("id") != task.id) or
                   (t.get("id", "").startswith(task.id + "-") or 
                    t.get("id", "").startswith(task.id + "_"))
            ]
            
            # If no subtasks found by ID pattern, check if task count increased
            if len(subtasks) == 0:
                if new_count > original_count:
                    # New tasks were added, assume they are subtasks
                    # Get tasks that weren't in the original list
                    subtasks = [t for t in updated_tasks if t.get("id") not in original_task_ids]
                
                if len(subtasks) == 0:
                    print("⚠️  No subtasks found. Agent may not have created them.")
                    print("   Checking if task was updated instead...")
                    
                    # Check if the original task was modified
                    original_task = next((t for t in updated_tasks if t.get("id") == task.id), None)
                    if original_task:
                        print(f"   Task {task.id} still exists in task list")
                        return False
                    else:
                        print(f"   Task {task.id} was removed but no subtasks found")
                        return False
            
            print(f"\n✅ Task broken down into {len(subtasks)} subtask(s):")
            for subtask in subtasks:
                print(f"   - {subtask['id']}: {subtask['title']} ({subtask.get('estimated_hours', 0)}h)")
            
            # Mark original task as failed (if it still exists)
            original_task = next((t for t in updated_tasks if t.get("id") == task.id), None)
            if original_task:
                original_task["status"] = "failed"
                original_task["description"] += f"\n\n[FAILED] {failure_reason}\n[REPLACED BY] {', '.join([s['id'] for s in subtasks])}"
                self._save_task_list(task_list)
            
            # Commit the changes
            self._commit_task_list()
            
            return True
        
        except Exception as e:
            print(f"\n❌ Error breaking down task: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def _build_breakdown_prompt(self, task: Task, failure_reason: str) -> str:
        """Build prompt for breaking down a task into subtasks."""
        # Load planning prompt template
        prompt_file = self.config.prompts_dir / "planning_prompt.md"
        
        if prompt_file.exists():
            with open(prompt_file, 'r') as f:
                template = f.read()
        else:
            template = "You are an AI planning agent. Break down complex tasks into smaller subtasks.\n\n"
        
        # Load current task list for context
        context = ""
        if self.config.task_list_path.exists():
            with open(self.config.task_list_path, 'r') as f:
                current_tasks = json.load(f)
                context += f"### Current task_list.json:\n```json\n{json.dumps(current_tasks, indent=2)}\n```\n\n"
        
        prompt = f"""{template}

## TASK BREAKDOWN REQUEST

A task has failed after multiple attempts and needs to be broken down into smaller, more manageable subtasks.

### Failed Task:

```json
{json.dumps(task.to_dict(), indent=2)}
```

### Failure Reason:

{failure_reason if failure_reason else "Task was too complex and could not be completed after multiple retry attempts."}

### Current Context:

{context}

### Your Task:

1. **Analyze the failed task** - Understand why it might have been too complex
2. **Break it down** into 3-5 smaller, focused subtasks
3. **Each subtask should be:**
   - 1-4 hours of work (preferably 2-3 hours)
   - Independently testable and completable
   - Have clear acceptance criteria
   - Specify files to modify
   
4. **Set dependencies:**
   - Subtasks should NOT depend on the original task ID (since it failed)
   - Set dependencies between subtasks if there's a logical order
   - If subtasks are independent, they can have empty dependencies
   
5. **Update task_list.json:**
   - Add the new subtasks to the tasks array
   - You can either:
     a) Mark the original task as "failed" and add subtasks
     b) Replace the original task with subtasks
   - Keep all other tasks unchanged

### Subtask Structure:

Each subtask should follow this format:

```json
{{
  "id": "{task.id}-SUB-1",
  "title": "Clear, focused title for this subtask",
  "priority": "{task.priority}",
  "status": "pending",
  "estimated_hours": 2.5,
  "description": "Detailed description of what this subtask accomplishes",
  "acceptance_criteria": [
    "Specific, testable criterion 1",
    "Specific, testable criterion 2",
    "Comprehensive test coverage"
  ],
  "dependencies": [],
  "files_to_modify": ["path/to/specific/file.dart"],
  "component": "{task.component}"
}}
```

### Guidelines:

- **Keep subtasks small**: 1-4 hours each, preferably 2-3 hours
- **Make them independent**: Each subtask should be completable on its own
- **Clear dependencies**: If subtasks depend on each other, set dependencies
- **Preserve priority**: Use the same priority as the original task
- **Specific files**: Each subtask should modify specific files, not all files from original task
- **Testable**: Each subtask should have clear acceptance criteria

### Example Breakdown:

If original task was: "Implement user authentication system (8 hours)"

Good breakdown:
- SUB-1: Create authentication models and DTOs (2h)
- SUB-2: Implement authentication service layer (2h)  
- SUB-3: Create authentication API endpoints (2h)
- SUB-4: Add authentication UI components (2h)

Now break down task `{task.id}` into subtasks and update the task_list.json file!
"""
        
        return prompt
    
    def _save_task_list(self, task_list: Dict[str, Any]) -> None:
        """Save task list to file."""
        with open(self.config.task_list_path, 'w') as f:
            json.dump(task_list, f, indent=2)


def main():
    """Main entry point."""
    if len(sys.argv) < 2:
        print("Usage: python planning_agent.py <component>")
        print("Components: web_v2, backend, common")
        sys.exit(1)
    
    component = sys.argv[1]
    
    # Load configuration
    config = AgentConfig.from_env(component)
    
    try:
        config.validate()
    except ValueError as e:
        print(f"❌ Configuration error: {e}")
        sys.exit(1)
    
    # Create planning agent and run
    agent = PlanningAgent(config)
    exit_code = agent.run()
    
    sys.exit(exit_code)


if __name__ == "__main__":
    main()