src/services/batch_service.py

"""
Module: services.batch_service
Description: Batch processing service integrating Celery and priority queue
Author: Anderson H. Silva
Date: 2025-01-25
License: Proprietary - All rights reserved
"""

from typing import Dict, Any, List, Optional, Callable
from datetime import datetime, timedelta
from enum import Enum
import asyncio

from pydantic import BaseModel, Field
from celery import group, chain, chord
from celery.result import AsyncResult

from src.core import get_logger
from src.infrastructure.queue.celery_app import celery_app, get_celery_app
from src.infrastructure.queue.priority_queue import (
    priority_queue,
    TaskPriority,
    TaskStatus,
    QueueStats
)
from src.infrastructure.queue.tasks import (
    run_investigation,
    analyze_contracts_batch,
    detect_anomalies_batch,
    analyze_patterns,
    generate_report,
    export_to_pdf,
    monitor_anomalies
)

logger = get_logger(__name__)


class BatchType(str, Enum):
    """Batch processing types."""
    INVESTIGATION = "investigation"
    ANALYSIS = "analysis"
    REPORT = "report"
    EXPORT = "export"
    MONITORING = "monitoring"


class BatchJobRequest(BaseModel):
    """Batch job request model."""
    batch_type: BatchType
    items: List[Dict[str, Any]]
    priority: TaskPriority = TaskPriority.NORMAL
    parallel: bool = True
    max_workers: int = 5
    callback_url: Optional[str] = None
    metadata: Dict[str, Any] = Field(default_factory=dict)


class BatchJobStatus(BaseModel):
    """Batch job status model."""
    job_id: str
    batch_type: BatchType
    total_items: int
    completed: int
    failed: int
    pending: int
    status: str
    started_at: datetime
    completed_at: Optional[datetime] = None
    duration_seconds: Optional[float] = None
    results: List[Dict[str, Any]] = Field(default_factory=list)


class BatchProcessingService:
    """Service for batch processing operations."""
    
    def __init__(self):
        """Initialize batch processing service."""
        self.celery_app = get_celery_app()
        self._active_jobs: Dict[str, BatchJobStatus] = {}
        self._job_results: Dict[str, List[AsyncResult]] = {}
        
        logger.info("batch_service_initialized")
    
    async def start(self):
        """Start batch processing service."""
        # Start priority queue
        await priority_queue.start()
        
        # Register handlers
        self._register_handlers()
        
        logger.info("batch_service_started")
    
    async def stop(self):
        """Stop batch processing service."""
        # Stop priority queue
        await priority_queue.stop()
        
        # Cancel active jobs
        for job_id, results in self._job_results.items():
            for result in results:
                if not result.ready():
                    result.revoke(terminate=True)
        
        logger.info("batch_service_stopped")
    
    def _register_handlers(self):
        """Register task handlers with priority queue."""
        # Investigation handler
        async def investigation_handler(payload: Dict[str, Any], metadata: Dict[str, Any]):
            result = run_investigation.delay(
                investigation_id=payload["investigation_id"],
                query=payload["query"],
                config=payload.get("config")
            )
            return result.id
        
        priority_queue.register_handler("investigation", investigation_handler)
        
        # Analysis handler
        async def analysis_handler(payload: Dict[str, Any], metadata: Dict[str, Any]):
            result = analyze_patterns.delay(
                data_type=payload["data_type"],
                time_range=payload["time_range"],
                pattern_types=payload.get("pattern_types"),
                min_confidence=payload.get("min_confidence", 0.7)
            )
            return result.id
        
        priority_queue.register_handler("analysis", analysis_handler)
    
    async def submit_batch_job(self, request: BatchJobRequest) -> BatchJobStatus:
        """
        Submit a batch job for processing.
        
        Args:
            request: Batch job request
            
        Returns:
            Batch job status
        """
        job_id = f"BATCH-{datetime.now().strftime('%Y%m%d%H%M%S')}"
        
        # Create job status
        job_status = BatchJobStatus(
            job_id=job_id,
            batch_type=request.batch_type,
            total_items=len(request.items),
            completed=0,
            failed=0,
            pending=len(request.items),
            status="submitted",
            started_at=datetime.now()
        )
        
        self._active_jobs[job_id] = job_status
        
        logger.info(
            "batch_job_submitted",
            job_id=job_id,
            batch_type=request.batch_type.value,
            items=len(request.items),
            priority=request.priority.name
        )
        
        # Create tasks based on batch type
        if request.batch_type == BatchType.INVESTIGATION:
            await self._process_investigation_batch(job_id, request)
        elif request.batch_type == BatchType.ANALYSIS:
            await self._process_analysis_batch(job_id, request)
        elif request.batch_type == BatchType.REPORT:
            await self._process_report_batch(job_id, request)
        elif request.batch_type == BatchType.EXPORT:
            await self._process_export_batch(job_id, request)
        elif request.batch_type == BatchType.MONITORING:
            await self._process_monitoring_batch(job_id, request)
        
        # Update status
        job_status.status = "processing"
        
        return job_status
    
    async def _process_investigation_batch(
        self,
        job_id: str,
        request: BatchJobRequest
    ):
        """Process investigation batch."""
        tasks = []
        
        for item in request.items:
            task = run_investigation.s(
                investigation_id=item.get("id", f"{job_id}-{len(tasks)}"),
                query=item["query"],
                config=item.get("config", {})
            )
            tasks.append(task)
        
        # Execute based on parallelism
        if request.parallel:
            job = group(tasks)
        else:
            job = chain(tasks)
        
        # Submit to Celery
        result = job.apply_async(
            priority=request.priority.value,
            link=self._create_callback_task(job_id, request.callback_url)
        )
        
        self._job_results[job_id] = [result]
    
    async def _process_analysis_batch(
        self,
        job_id: str,
        request: BatchJobRequest
    ):
        """Process analysis batch."""
        tasks = []
        
        for item in request.items:
            if item.get("type") == "contracts":
                task = analyze_contracts_batch.s(
                    contract_ids=item["contract_ids"],
                    analysis_type=item.get("analysis_type", "anomaly"),
                    threshold=item.get("threshold", 0.7)
                )
            elif item.get("type") == "patterns":
                task = analyze_patterns.s(
                    data_type=item["data_type"],
                    time_range=item["time_range"],
                    pattern_types=item.get("pattern_types"),
                    min_confidence=item.get("min_confidence", 0.7)
                )
            else:
                continue
            
            tasks.append(task)
        
        # Execute in parallel
        job = group(tasks)
        result = job.apply_async(
            priority=request.priority.value,
            link=self._create_callback_task(job_id, request.callback_url)
        )
        
        self._job_results[job_id] = [result]
    
    async def _process_report_batch(
        self,
        job_id: str,
        request: BatchJobRequest
    ):
        """Process report batch."""
        tasks = []
        
        for item in request.items:
            task = generate_report.s(
                report_id=item.get("id", f"{job_id}-{len(tasks)}"),
                report_type=item["report_type"],
                investigation_ids=item["investigation_ids"],
                config=item.get("config", {})
            )
            tasks.append(task)
        
        # Generate reports in parallel
        job = group(tasks)
        result = job.apply_async(
            priority=request.priority.value,
            link=self._create_callback_task(job_id, request.callback_url)
        )
        
        self._job_results[job_id] = [result]
    
    async def _process_export_batch(
        self,
        job_id: str,
        request: BatchJobRequest
    ):
        """Process export batch."""
        tasks = []
        
        for item in request.items:
            task = export_to_pdf.s(
                content_type=item["content_type"],
                content_id=item["content_id"],
                options=item.get("options", {})
            )
            tasks.append(task)
        
        # Export in parallel with limited workers
        job = group(tasks)
        result = job.apply_async(
            priority=request.priority.value,
            link=self._create_callback_task(job_id, request.callback_url),
            queue="normal"
        )
        
        self._job_results[job_id] = [result]
    
    async def _process_monitoring_batch(
        self,
        job_id: str,
        request: BatchJobRequest
    ):
        """Process monitoring batch."""
        tasks = []
        
        for item in request.items:
            task = monitor_anomalies.s(
                monitoring_config=item["config"],
                alert_threshold=item.get("threshold", 0.8)
            )
            tasks.append(task)
        
        # Run monitoring tasks
        job = group(tasks)
        result = job.apply_async(
            priority=request.priority.value,
            link=self._create_callback_task(job_id, request.callback_url)
        )
        
        self._job_results[job_id] = [result]
    
    def _create_callback_task(self, job_id: str, callback_url: Optional[str]):
        """Create callback task for job completion."""
        if not callback_url:
            return None
        
        @celery_app.task
        def batch_completion_callback(results):
            # Update job status
            job_status = self._active_jobs.get(job_id)
            if job_status:
                job_status.completed_at = datetime.now()
                job_status.duration_seconds = (
                    job_status.completed_at - job_status.started_at
                ).total_seconds()
                job_status.status = "completed"
                job_status.results = results
            
            # Send callback
            import httpx
            with httpx.Client() as client:
                client.post(
                    callback_url,
                    json={
                        "job_id": job_id,
                        "status": "completed",
                        "results": results,
                        "completed_at": datetime.now().isoformat()
                    },
                    timeout=30.0
                )
        
        return batch_completion_callback.s()
    
    async def get_job_status(self, job_id: str) -> Optional[BatchJobStatus]:
        """
        Get batch job status.
        
        Args:
            job_id: Job ID
            
        Returns:
            Job status or None
        """
        job_status = self._active_jobs.get(job_id)
        if not job_status:
            return None
        
        # Update status from Celery results
        if job_id in self._job_results:
            results = self._job_results[job_id]
            completed = 0
            failed = 0
            
            for result in results:
                if result.ready():
                    if result.successful():
                        completed += 1
                    else:
                        failed += 1
            
            job_status.completed = completed
            job_status.failed = failed
            job_status.pending = job_status.total_items - completed - failed
            
            if job_status.pending == 0:
                job_status.status = "completed" if failed == 0 else "completed_with_errors"
                if not job_status.completed_at:
                    job_status.completed_at = datetime.now()
                    job_status.duration_seconds = (
                        job_status.completed_at - job_status.started_at
                    ).total_seconds()
        
        return job_status
    
    async def cancel_job(self, job_id: str) -> bool:
        """
        Cancel a batch job.
        
        Args:
            job_id: Job ID
            
        Returns:
            True if cancelled
        """
        if job_id not in self._job_results:
            return False
        
        # Revoke Celery tasks
        for result in self._job_results[job_id]:
            if not result.ready():
                result.revoke(terminate=True)
        
        # Update status
        job_status = self._active_jobs.get(job_id)
        if job_status:
            job_status.status = "cancelled"
            job_status.completed_at = datetime.now()
            job_status.duration_seconds = (
                job_status.completed_at - job_status.started_at
            ).total_seconds()
        
        logger.info("batch_job_cancelled", job_id=job_id)
        
        return True
    
    async def get_queue_stats(self) -> QueueStats:
        """Get queue statistics."""
        return await priority_queue.get_stats()
    
    async def cleanup_old_jobs(self, days: int = 7):
        """Clean up old completed jobs."""
        cutoff_time = datetime.now() - timedelta(days=days)
        
        jobs_to_remove = []
        for job_id, job_status in self._active_jobs.items():
            if (job_status.completed_at and 
                job_status.completed_at < cutoff_time):
                jobs_to_remove.append(job_id)
        
        for job_id in jobs_to_remove:
            del self._active_jobs[job_id]
            if job_id in self._job_results:
                del self._job_results[job_id]
        
        logger.info(
            "old_jobs_cleaned",
            removed=len(jobs_to_remove),
            remaining=len(self._active_jobs)
        )


# Global batch service instance
batch_service = BatchProcessingService()