mnemo_cards/tools/agent/task_manager.py

"""
Task and state management for AI agents.
"""
import json
import fcntl
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, asdict


@dataclass
class Task:
    """Represents a single task."""
    id: str
    title: str
    priority: str  # 'high', 'medium', 'low'
    status: str  # 'pending', 'in_progress', 'completed', 'failed', 'skipped'
    estimated_hours: float
    description: str
    acceptance_criteria: List[str]
    dependencies: List[str]
    files_to_modify: List[str]
    component: str = "web_v2"
    
    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Task":
        return cls(**data)


@dataclass
class AgentState:
    """Represents the current state of an agent."""
    component: str
    current_task_id: Optional[str]
    iteration_count: int
    max_iterations: int
    started_at: Optional[str]
    last_commit: Optional[str]
    retry_count: int
    max_retries: int
    status: str  # 'idle', 'in_progress', 'completed', 'error'
    errors: List[str]
    completed_tasks: List[str]
    skipped_tasks: List[str]
    
    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "AgentState":
        return cls(**data)


class TaskManager:
    """Manages tasks and agent state."""
    
    def __init__(self, task_list_path: Path, agent_state_path: Path):
        self.task_list_path = task_list_path
        self.agent_state_path = agent_state_path
    
    def _read_json_file(self, path: Path) -> Dict[str, Any]:
        """Read and parse JSON file with file locking."""
        with open(path, 'r') as f:
            fcntl.flock(f.fileno(), fcntl.LOCK_SH)
            try:
                return json.load(f)
            finally:
                fcntl.flock(f.fileno(), fcntl.LOCK_UN)
    
    def _write_json_file(self, path: Path, data: Dict[str, Any]) -> None:
        """Write JSON file with file locking."""
        with open(path, 'w') as f:
            fcntl.flock(f.fileno(), fcntl.LOCK_EX)
            try:
                json.dump(data, f, indent=2)
            finally:
                fcntl.flock(f.fileno(), fcntl.LOCK_UN)
    
    def load_tasks(self) -> List[Task]:
        """Load tasks from task list file."""
        data = self._read_json_file(self.task_list_path)
        return [Task.from_dict(task_data) for task_data in data.get("tasks", [])]
    
    def save_tasks(self, tasks: List[Task]) -> None:
        """Save tasks to task list file."""
        data = self._read_json_file(self.task_list_path)
        data["tasks"] = [task.to_dict() for task in tasks]
        self._write_json_file(self.task_list_path, data)
    
    def load_state(self) -> AgentState:
        """Load agent state from file."""
        data = self._read_json_file(self.agent_state_path)
        return AgentState.from_dict(data)
    
    def save_state(self, state: AgentState) -> None:
        """Save agent state to file."""
        self._write_json_file(self.agent_state_path, state.to_dict())
    
    def get_next_task(self) -> Optional[Task]:
        """Get next pending task by priority."""
        tasks = self.load_tasks()
        state = self.load_state()
        
        # Filter pending tasks
        pending_tasks = [
            task for task in tasks
            if task.status == "pending" and task.id not in state.skipped_tasks
        ]
        
        if not pending_tasks:
            return None
        
        # Check dependencies
        completed_task_ids = set(state.completed_tasks)
        
        def can_start_task(task: Task) -> bool:
            """Check if task dependencies are satisfied."""
            if not task.dependencies:
                return True
            
            # Parse dependencies (format: "component:TASK-ID" or "TASK-ID")
            for dep in task.dependencies:
                if ":" in dep:
                    dep_component, dep_id = dep.split(":", 1)
                    # TODO: Check other component's state
                    # For now, only check current component
                    if dep_component == state.component and dep_id not in completed_task_ids:
                        return False
                else:
                    if dep not in completed_task_ids:
                        return False
            
            return True
        
        # Filter by dependencies
        ready_tasks = [task for task in pending_tasks if can_start_task(task)]
        
        if not ready_tasks:
            return None
        
        # Sort by priority: high > medium > low
        priority_order = {"high": 0, "medium": 1, "low": 2}
        ready_tasks.sort(key=lambda t: priority_order.get(t.priority, 3))
        
        return ready_tasks[0]
    
    def update_task_status(self, task_id: str, status: str) -> None:
        """Update task status."""
        tasks = self.load_tasks()
        for task in tasks:
            if task.id == task_id:
                task.status = status
                break
        self.save_tasks(tasks)
    
    def mark_task_completed(self, task_id: str) -> None:
        """Mark task as completed."""
        self.update_task_status(task_id, "completed")
        
        state = self.load_state()
        if task_id not in state.completed_tasks:
            state.completed_tasks.append(task_id)
        state.current_task_id = None
        state.retry_count = 0
        self.save_state(state)
    
    def mark_task_failed(self, task_id: str, error: str) -> None:
        """Mark task as failed."""
        self.update_task_status(task_id, "failed")
        
        state = self.load_state()
        state.errors.append(f"{task_id}: {error}")
        state.current_task_id = None
        state.status = "idle"  # Reset status to allow agent to continue
        state.retry_count = 0  # Reset retry count for next task
        self.save_state(state)
    
    def mark_task_skipped(self, task_id: str, reason: str) -> None:
        """Mark task as skipped."""
        self.update_task_status(task_id, "skipped")
        
        state = self.load_state()
        if task_id not in state.skipped_tasks:
            state.skipped_tasks.append(task_id)
        state.errors.append(f"{task_id} skipped: {reason}")
        state.current_task_id = None
        state.status = "idle"  # Reset status to allow agent to continue
        state.retry_count = 0  # Reset retry count for next task
        self.save_state(state)
    
    def start_task(self, task: Task) -> None:
        """Mark task as started."""
        self.update_task_status(task.id, "in_progress")
        
        state = self.load_state()
        state.current_task_id = task.id
        state.status = "in_progress"
        if not state.started_at:
            state.started_at = datetime.now(timezone.utc).isoformat()
        state.iteration_count += 1
        self.save_state(state)
    
    def increment_retry(self) -> int:
        """Increment retry count and return new value."""
        state = self.load_state()
        state.retry_count += 1
        self.save_state(state)
        return state.retry_count
    
    def is_agent_running(self) -> bool:
        """Check if agent is currently running."""
        state = self.load_state()
        
        if state.status != "in_progress":
            return False
        
        # Check if started too long ago (stale lock)
        if state.started_at:
            started = datetime.fromisoformat(state.started_at)
            now = datetime.now(timezone.utc)
            hours_elapsed = (now - started).total_seconds() / 3600
            
            if hours_elapsed > 2:  # Consider stale after 2 hours
                return False
        
        return True
    
    def reset_state(self) -> None:
        """Reset agent state to idle."""
        state = self.load_state()
        state.status = "idle"
        state.current_task_id = None
        state.started_at = None
        state.iteration_count = 0
        state.retry_count = 0
        self.save_state(state)
    
    def all_tasks_completed(self) -> bool:
        """Check if all tasks are completed."""
        tasks = self.load_tasks()
        return all(task.status in ["completed", "skipped"] for task in tasks)


class GlobalLock:
    """Manages global lock for shared resource access."""
    
    def __init__(self, lock_file: Path):
        self.lock_file = lock_file
    
    def _read_lock(self) -> Dict[str, Any]:
        """Read lock file."""
        with open(self.lock_file, 'r') as f:
            fcntl.flock(f.fileno(), fcntl.LOCK_SH)
            try:
                return json.load(f)
            finally:
                fcntl.flock(f.fileno(), fcntl.LOCK_UN)
    
    def _write_lock(self, data: Dict[str, Any]) -> None:
        """Write lock file."""
        with open(self.lock_file, 'w') as f:
            fcntl.flock(f.fileno(), fcntl.LOCK_EX)
            try:
                json.dump(data, f, indent=2)
            finally:
                fcntl.flock(f.fileno(), fcntl.LOCK_UN)
    
    def acquire(self, component: str, reason: str = "Working on shared resources") -> bool:
        """Acquire global lock."""
        lock_data = self._read_lock()
        
        if lock_data.get("locked"):
            return False
        
        lock_data["locked"] = True
        lock_data["locked_by"] = component
        lock_data["locked_at"] = datetime.now(timezone.utc).isoformat()
        lock_data["reason"] = reason
        
        self._write_lock(lock_data)
        return True
    
    def release(self, component: str) -> None:
        """Release global lock."""
        lock_data = self._read_lock()
        
        if lock_data.get("locked_by") == component:
            lock_data["locked"] = False
            lock_data["locked_by"] = None
            lock_data["locked_at"] = None
            lock_data["reason"] = None
            self._write_lock(lock_data)
    
    def is_locked(self) -> bool:
        """Check if global lock is held."""
        lock_data = self._read_lock()
        return lock_data.get("locked", False)