"""
Batch processing system for Trax platform.

Provides async batch processing with configurable workers, progress tracking,
error recovery, and resource monitoring.
"""

import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Protocol, Set
from enum import Enum
import psutil

from .media_types import ProgressCallback
from .transcription_service import TranscriptionService, TranscriptionConfig
from .enhancement import EnhancementServiceProtocol
from .media_service import MediaService

logger = logging.getLogger(__name__)


class TaskType(Enum):
    """Types of tasks that can be processed in batch."""
    TRANSCRIBE = "transcribe"
    ENHANCE = "enhance"
    YOUTUBE = "youtube"
    DOWNLOAD = "download"
    PREPROCESS = "preprocess"


@dataclass
class BatchTask:
    """Represents a single task in the batch processing queue."""
    id: str
    task_type: TaskType
    data: Dict[str, Any]
    priority: int = 0
    retry_count: int = 0
    max_retries: int = 3
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    error: Optional[str] = None
    result: Optional[Dict[str, Any]] = None


@dataclass
class BatchProgress:
    """Tracks progress of batch processing."""
    total_tasks: int
    completed_tasks: int = 0
    failed_tasks: int = 0
    in_progress_tasks: int = 0
    queued_tasks: int = 0
    start_time: Optional[datetime] = None
    estimated_completion: Optional[datetime] = None
    current_worker_count: int = 0
    memory_usage_mb: float = 0.0
    cpu_usage_percent: float = 0.0
    
    @property
    def success_rate(self) -> float:
        return (self.completed_tasks / self.total_tasks * 100) if self.total_tasks > 0 else 0.0
    
    @property
    def failure_rate(self) -> float:
        return (self.failed_tasks / self.total_tasks * 100) if self.total_tasks > 0 else 0.0
    
    @property
    def elapsed_time(self) -> Optional[float]:
        if self.start_time is None:
            return None
        return (datetime.now(timezone.utc) - self.start_time).total_seconds()


@dataclass
class BatchResult:
    """Result of batch processing operation."""
    success_count: int
    failure_count: int
    total_count: int
    results: List[Dict[str, Any]]
    failures: List[Dict[str, Any]]
    processing_time: float
    memory_peak_mb: float
    cpu_peak_percent: float
    quality_metrics: Dict[str, float]
    
    @property
    def success_rate(self) -> float:
        return (self.success_count / self.total_count * 100) if self.total_count > 0 else 0.0


class BatchProcessorProtocol(Protocol):
    """Protocol for batch processing services."""
    async def add_task(self, task_type: TaskType, data: Dict[str, Any], priority: int = 0) -> str: ...
    async def start(self, progress_callback: Optional[Callable[[BatchProgress], None]] = None) -> BatchResult: ...
    async def pause(self) -> None: ...
    async def resume(self) -> None: ...
    async def stop(self) -> None: ...
    def get_progress(self) -> BatchProgress: ...


class BatchProcessor:
    """Async batch processor with configurable worker pool."""
    
    def __init__(
        self,
        max_workers: int = 8,
        queue_size: int = 1000,
        progress_interval: float = 5.0,
        memory_limit_mb: float = 2048.0,
        cpu_limit_percent: float = 90.0
    ):
        self.max_workers = max_workers
        self.progress_interval = progress_interval
        self.memory_limit_mb = memory_limit_mb
        self.cpu_limit_percent = cpu_limit_percent
        
        # State
        self.running = False
        self.paused = False
        self.stopped = False
        
        # Queues and storage
        self.task_queue: asyncio.PriorityQueue = asyncio.PriorityQueue(maxsize=queue_size)
        self.completed_tasks: List[BatchTask] = []
        self.failed_tasks: List[BatchTask] = []
        self.active_tasks: Set[str] = set()
        
        # Progress tracking
        self.progress = BatchProgress(total_tasks=0)
        self.progress_callback: Optional[Callable[[BatchProgress], None]] = None
        self.progress_task: Optional[asyncio.Task] = None
        
        # Resource monitoring
        self.memory_peak_mb = 0.0
        self.cpu_peak_percent = 0.0
        
        # Services
        self.transcription_service: Optional[TranscriptionService] = None
        self.enhancement_service: Optional[EnhancementService] = None
        self.media_service: Optional[MediaService] = None
        
        # Workers
        self.workers: List[asyncio.Task] = []
        self.semaphore = asyncio.Semaphore(max_workers)
        self.task_counter = 0
        
        logger.info(f"BatchProcessor initialized with {max_workers} workers")
    
    async def add_task(self, task_type: TaskType, data: Dict[str, Any], priority: int = 0) -> str:
        """Add a task to the processing queue."""
        self.task_counter += 1
        task_id = f"task_{self.task_counter}_{task_type.value}"
        
        task = BatchTask(id=task_id, task_type=task_type, data=data, priority=priority)
        await self.task_queue.put((priority, task))
        
        self.progress.queued_tasks += 1
        self.progress.total_tasks += 1
        
        logger.debug(f"Added task {task_id} to queue (priority: {priority})")
        return task_id
    
    async def _initialize_services(self) -> None:
        """Initialize required services if not already done."""
        if self.transcription_service is None:
            from .transcription_service import create_transcription_service
            self.transcription_service = create_transcription_service()
            await self.transcription_service.initialize()
        
        if self.enhancement_service is None:
            from .enhancement.enhancement_service import create_enhancement_service
            self.enhancement_service = create_enhancement_service()
        
        if self.media_service is None:
            from .media_service import create_media_service
            from ..repositories.media_repository import create_media_repository
            media_repo = create_media_repository()
            self.media_service = create_media_service(media_repo)
    
    async def _process_task(self, task: BatchTask) -> Dict[str, Any]:
        """Process a single task based on its type."""
        task.started_at = datetime.now(timezone.utc)
        self.active_tasks.add(task.id)
        self.progress.in_progress_tasks += 1
        self.progress.queued_tasks -= 1
        
        try:
            logger.info(f"Processing task {task.id} ({task.task_type.value})")
            
            if task.task_type == TaskType.TRANSCRIBE:
                result = await self._process_transcription(task)
            elif task.task_type == TaskType.ENHANCE:
                result = await self._process_enhancement(task)
            elif task.task_type == TaskType.YOUTUBE:
                result = await self._process_youtube(task)
            elif task.task_type == TaskType.DOWNLOAD:
                result = await self._process_download(task)
            elif task.task_type == TaskType.PREPROCESS:
                result = await self._process_preprocessing(task)
            else:
                raise ValueError(f"Unknown task type: {task.task_type}")
            
            task.completed_at = datetime.now(timezone.utc)
            task.result = result
            self.completed_tasks.append(task)
            self.progress.completed_tasks += 1
            
            logger.info(f"Task {task.id} completed successfully")
            return result
            
        except Exception as e:
            task.error = str(e)
            task.completed_at = datetime.now(timezone.utc)
            
            if task.retry_count < task.max_retries:
                task.retry_count += 1
                task.started_at = None
                task.completed_at = None
                task.error = None
                await self.task_queue.put((task.priority + 1, task))
                self.progress.queued_tasks += 1
                logger.warning(f"Task {task.id} failed, retrying ({task.retry_count}/{task.max_retries})")
                return {"status": "retrying", "retry_count": task.retry_count}
            else:
                self.failed_tasks.append(task)
                self.progress.failed_tasks += 1
                logger.error(f"Task {task.id} failed permanently after {task.max_retries} retries: {e}")
                return {"status": "failed", "error": str(e)}
        
        finally:
            self.active_tasks.discard(task.id)
            self.progress.in_progress_tasks -= 1
    
    async def _process_transcription(self, task: BatchTask) -> Dict[str, Any]:
        """Process transcription task."""
        if self.transcription_service is None:
            raise RuntimeError("Transcription service not initialized")
        
        file_path = Path(task.data["file_path"])
        config = TranscriptionConfig(**task.data.get("config", {}))
        result = await self.transcription_service.transcribe_file(file_path, config)
        
        return {
            "status": "completed",
            "file_path": str(file_path),
            "transcript": result.text_content,
            "segments": result.segments,
            "accuracy": result.accuracy,
            "processing_time": result.processing_time,
            "quality_warnings": result.quality_warnings
        }
    
    async def _process_enhancement(self, task: BatchTask) -> Dict[str, Any]:
        """Process enhancement task."""
        if self.enhancement_service is None:
            raise RuntimeError("Enhancement service not initialized")
        
        transcript_id = task.data["transcript_id"]
        result = await self.enhancement_service.enhance_transcript(transcript_id)
        
        return {
            "status": "completed",
            "transcript_id": transcript_id,
            "enhanced_content": result.enhanced_content,
            "accuracy_improvement": result.accuracy_improvement,
            "processing_time": result.processing_time
        }
    
    async def _process_youtube(self, task: BatchTask) -> Dict[str, Any]:
        """Process YouTube metadata extraction task."""
        url = task.data["url"]
        return {"status": "completed", "url": url, "metadata": {"title": "Placeholder", "duration": 0}}
    
    async def _process_download(self, task: BatchTask) -> Dict[str, Any]:
        """Process media download task."""
        if self.media_service is None:
            raise RuntimeError("Media service not initialized")
        
        url = task.data["url"]
        result = await self.media_service.download_media(url)
        
        return {
            "status": "completed",
            "url": url,
            "file_path": str(result.file_path),
            "file_size": result.file_size,
            "duration": result.duration
        }
    
    async def _process_preprocessing(self, task: BatchTask) -> Dict[str, Any]:
        """Process media preprocessing task."""
        if self.media_service is None:
            raise RuntimeError("Media service not initialized")
        
        file_path = Path(task.data["file_path"])
        result = await self.media_service.preprocess_media(file_path)
        
        return {
            "status": "completed",
            "file_path": str(file_path),
            "processed_path": str(result.processed_path),
            "processing_time": result.processing_time
        }
    
    async def _worker(self, worker_id: int) -> None:
        """Worker function that processes tasks from the queue."""
        logger.debug(f"Worker {worker_id} started")
        
        while self.running and not self.stopped:
            try:
                if self.paused:
                    await asyncio.sleep(1.0)
                    continue
                
                try:
                    async with self.semaphore:
                        priority, task = await asyncio.wait_for(self.task_queue.get(), timeout=1.0)
                except asyncio.TimeoutError:
                    if self.task_queue.empty():
                        break
                    continue
                
                await self._process_task(task)
                self.task_queue.task_done()
                
            except Exception as e:
                logger.error(f"Worker {worker_id} error: {e}")
                await asyncio.sleep(1.0)
        
        logger.debug(f"Worker {worker_id} stopped")
    
    async def _progress_monitor(self) -> None:
        """Monitor and report progress at regular intervals."""
        while self.running and not self.stopped:
            try:
                process = psutil.Process()
                memory_mb = process.memory_info().rss / 1024 / 1024
                cpu_percent = process.cpu_percent()
                
                self.progress.memory_usage_mb = memory_mb
                self.progress.cpu_usage_percent = cpu_percent
                self.progress.current_worker_count = len(self.active_tasks)
                
                self.memory_peak_mb = max(self.memory_peak_mb, memory_mb)
                self.cpu_peak_percent = max(self.cpu_peak_percent, cpu_percent)
                
                if self.progress_callback:
                    self.progress_callback(self.progress)
                
                await asyncio.sleep(self.progress_interval)
                
            except Exception as e:
                logger.error(f"Progress monitor error: {e}")
                await asyncio.sleep(self.progress_interval)
    
    async def start(self, progress_callback: Optional[Callable[[BatchProgress], None]] = None) -> BatchResult:
        """Start batch processing."""
        if self.running:
            raise RuntimeError("Batch processor is already running")
        
        logger.info("Starting batch processing")
        
        self.running = True
        self.paused = False
        self.stopped = False
        self.progress.start_time = datetime.now(timezone.utc)
        self.progress_callback = progress_callback
        
        await self._initialize_services()
        self.progress_task = asyncio.create_task(self._progress_monitor())
        
        self.workers = [asyncio.create_task(self._worker(i)) for i in range(self.max_workers)]
        
        try:
            await self.task_queue.join()
            self.running = False
            await asyncio.gather(*self.workers, return_exceptions=True)
            
            if self.progress_task:
                self.progress_task.cancel()
                try:
                    await self.progress_task
                except asyncio.CancelledError:
                    pass
            
            processing_time = self.progress.elapsed_time or 0.0
            
            quality_metrics = {}
            if self.completed_tasks:
                transcription_tasks = [t for t in self.completed_tasks if t.task_type == TaskType.TRANSCRIBE]
                if transcription_tasks:
                    avg_accuracy = sum(t.result.get("accuracy", 0) for t in transcription_tasks) / len(transcription_tasks)
                    quality_metrics["avg_transcription_accuracy"] = avg_accuracy
                
                enhancement_tasks = [t for t in self.completed_tasks if t.task_type == TaskType.ENHANCE]
                if enhancement_tasks:
                    avg_improvement = sum(t.result.get("accuracy_improvement", 0) for t in enhancement_tasks) / len(enhancement_tasks)
                    quality_metrics["avg_enhancement_improvement"] = avg_improvement
            
            result = BatchResult(
                success_count=len(self.completed_tasks),
                failure_count=len(self.failed_tasks),
                total_count=self.progress.total_tasks,
                results=[t.result for t in self.completed_tasks if t.result],
                failures=[{"task_id": t.id, "error": t.error} for t in self.failed_tasks],
                processing_time=processing_time,
                memory_peak_mb=self.memory_peak_mb,
                cpu_peak_percent=self.cpu_peak_percent,
                quality_metrics=quality_metrics
            )
            
            logger.info(f"Batch processing completed: {result.success_count} successful, {result.failure_count} failed")
            return result
            
        except Exception as e:
            logger.error(f"Batch processing error: {e}")
            raise
        finally:
            self.running = False
    
    async def pause(self) -> None:
        """Pause batch processing."""
        if self.running:
            self.paused = True
            logger.info("Batch processing paused")
    
    async def resume(self) -> None:
        """Resume batch processing."""
        if self.running:
            self.paused = False
            logger.info("Batch processing resumed")
    
    async def stop(self) -> None:
        """Stop batch processing."""
        self.stopped = True
        self.running = False
        logger.info("Batch processing stopped")
    
    def get_progress(self) -> BatchProgress:
        """Get current progress."""
        return self.progress


def create_batch_processor(
    max_workers: int = 8,
    queue_size: int = 1000,
    progress_interval: float = 5.0,
    memory_limit_mb: float = 2048.0,
    cpu_limit_percent: float = 90.0
) -> BatchProcessor:
    """Create a new batch processor instance."""
    return BatchProcessor(
        max_workers=max_workers,
        queue_size=queue_size,
        progress_interval=progress_interval,
        memory_limit_mb=memory_limit_mb,
        cpu_limit_percent=cpu_limit_percent
    )
