"""
Task manager for AgentChain.
"""

import json
import logging
from pathlib import Path
from typing import Optional, List, Dict, Any

from .models import Task, Plan, TaskStatus


logger = logging.getLogger(__name__)


class TaskManager:
    """
    Manages task execution and persistence.
    """
    
    def __init__(self, workspace_dir: str = "."):
        self.workspace_dir = Path(workspace_dir)
        self.data_dir = self.workspace_dir / ".llama2a"
        self.data_dir.mkdir(parents=True, exist_ok=True)
        
        self.current_plan: Optional[Plan] = None
    
    def create_plan(self, name: str, description: str = "") -> Plan:
        """Create a new execution plan."""
        self.current_plan = Plan(name=name, description=description)
        return self.current_plan
    
    def load_plan(self, plan_data: Dict[str, Any]) -> Plan:
        """Load a plan from dictionary data."""
        self.current_plan = Plan.from_dict(plan_data)
        return self.current_plan
    
    def add_task(self, task: Task) -> None:
        """Add a task to the current plan."""
        if not self.current_plan:
            self.create_plan("Default Plan")
        self.current_plan.add_task(task)
    
    def get_next_task(self) -> Optional[Task]:
        """Get the next task to execute."""
        if not self.current_plan:
            return None
        return self.current_plan.get_next_task()
    
    def update_task_status(
        self,
        task_id: str,
        status: TaskStatus,
        output: Optional[str] = None,
        error: Optional[str] = None
    ) -> None:
        """Update a task's status."""
        if not self.current_plan:
            return
        
        task = self.current_plan.get_task(task_id)
        if task:
            task.status = status
            if output:
                task.output = output
            if error:
                task.error = error
            
            logger.info(f"Task {task_id} status: {status.value}")
            self.save_plan()
    
    def is_complete(self) -> bool:
        """Check if all tasks are complete."""
        return self.current_plan.is_complete() if self.current_plan else True
    
    def get_progress(self) -> Dict[str, Any]:
        """Get current progress."""
        if not self.current_plan:
            return {"total": 0, "completed": 0, "percent": 0}
        return self.current_plan.get_progress()
    
    def get_failed_tasks(self) -> List[Task]:
        """Get all failed tasks."""
        if not self.current_plan:
            return []
        return [t for t in self.current_plan.tasks if t.status == TaskStatus.FAILED]
    
    def reset_failed_tasks(self) -> int:
        """Reset failed tasks for retry."""
        count = 0
        for task in self.get_failed_tasks():
            if task.can_retry():
                task.status = TaskStatus.PENDING
                count += 1
        return count
    
    def save_plan(self, filename: str = "plan.json") -> None:
        """Save current plan to file."""
        if not self.current_plan:
            return
        
        path = self.data_dir / filename
        with open(path, "w") as f:
            json.dump(self.current_plan.to_dict(), f, indent=2, default=str)
        logger.debug(f"Saved plan to {path}")
    
    def load_plan_from_file(self, filename: str = "plan.json") -> Optional[Plan]:
        """Load plan from file."""
        path = self.data_dir / filename
        if not path.exists():
            return None
        
        with open(path) as f:
            data = json.load(f)
        
        self.current_plan = Plan.from_dict(data)
        return self.current_plan
