"""
Progress Tracking System

Manages progress updates for long-running enhancement tasks
Supports Server-Sent Events (SSE) for real-time frontend updates
"""
import asyncio
import logging
from typing import Dict, Optional
from datetime import datetime


logger = logging.getLogger(__name__)


class ProgressTracker:
    """
    Tracks progress for enhancement operations
    Allows multiple clients to subscribe to progress updates
    """
    
    def __init__(self):
        self.tasks: Dict[str, Dict] = {}
        self.subscribers: Dict[str, list] = {}
    
    def create_task(self, task_id: str, total_steps: int = 100):
        """Initialize a new task for progress tracking"""
        self.tasks[task_id] = {
            "id": task_id,
            "status": "initializing",
            "progress": 0,
            "total_steps": total_steps,
            "current_step": 0,
            "message": "Starting enhancement...",
            "started_at": datetime.now().isoformat(),
            "error": None
        }
        self.subscribers[task_id] = []
        logger.info(f"Created progress tracker for task {task_id}")
    
    def update_progress(
        self, 
        task_id: str, 
        current_step: int, 
        message: str,
        status: str = "processing"
    ):
        """Update progress for a task (synchronous)"""
        if task_id not in self.tasks:
            logger.warning(f"Task {task_id} not found for progress update")
            return
        
        task = self.tasks[task_id]
        task["current_step"] = current_step
        task["progress"] = int((current_step / task["total_steps"]) * 100)
        task["message"] = message
        task["status"] = status
        
        logger.info(f"Task {task_id}: {task['progress']}% - {message}")
    
    def complete_task(self, task_id: str, result_url: Optional[str] = None):
        """Mark a task as completed"""
        if task_id not in self.tasks:
            return
        
        self.tasks[task_id].update({
            "status": "completed",
            "progress": 100,
            "message": "Enhancement completed!",
            "result_url": result_url,
            "completed_at": datetime.now().isoformat()
        })
        
        # Notify subscribers if event loop is running
        try:
            asyncio.create_task(self._notify_subscribers(task_id))
        except RuntimeError:
            # No event loop running (called from thread)
            pass
        
        logger.info(f"Task {task_id} completed")
    
    def fail_task(self, task_id: str, error: str):
        """Mark a task as failed"""
        if task_id not in self.tasks:
            return
        
        self.tasks[task_id].update({
            "status": "failed",
            "message": f"Error: {error}",
            "error": error,
            "failed_at": datetime.now().isoformat()
        })
        
        # Notify subscribers if event loop is running
        try:
            asyncio.create_task(self._notify_subscribers(task_id))
        except RuntimeError:
            # No event loop running (called from thread)
            pass
        
        logger.error(f"Task {task_id} failed: {error}")
    
    def get_task_status(self, task_id: str) -> Optional[Dict]:
        """Get current status of a task"""
        return self.tasks.get(task_id)
    
    def subscribe(self, task_id: str, queue: asyncio.Queue):
        """Subscribe to progress updates for a task"""
        if task_id not in self.subscribers:
            self.subscribers[task_id] = []
        self.subscribers[task_id].append(queue)
        logger.debug(f"New subscriber for task {task_id}")
    
    def unsubscribe(self, task_id: str, queue: asyncio.Queue):
        """Unsubscribe from progress updates"""
        if task_id in self.subscribers and queue in self.subscribers[task_id]:
            self.subscribers[task_id].remove(queue)
    
    async def _notify_subscribers(self, task_id: str):
        """Send current task status to all subscribers"""
        if task_id not in self.subscribers:
            return
        
        task_data = self.tasks.get(task_id)
        if not task_data:
            return
        
        # Send to all subscribers
        for queue in self.subscribers[task_id]:
            try:
                await queue.put(task_data.copy())
            except Exception as e:
                logger.error(f"Error notifying subscriber: {e}")
    
    def cleanup_task(self, task_id: str):
        """Remove task data after completion"""
        if task_id in self.tasks:
            del self.tasks[task_id]
        if task_id in self.subscribers:
            del self.subscribers[task_id]
        logger.debug(f"Cleaned up task {task_id}")


# Global progress tracker instance
progress_tracker = ProgressTracker()