backend/queue/producer.py

"""
Job Producer for Evaluation Queue

Provides job submission functionality for the evaluation queue.
Handles job creation, validation, and queueing.
"""

import hashlib
import json
import uuid
from datetime import datetime
from typing import Optional

from backend.core.config import settings
from backend.logging.logger import get_logger

from .job_schema import (
    JobStatus,
    JobType,
    JobPriority,
    EvaluationJob,
    JobSubmissionRequest,
)
from .status_tracker import get_status_tracker


logger = get_logger("queue.producer", component="queue")


class JobProducer:
    """
    Produces and submits evaluation jobs to the queue.
    
    Responsibilities:
    - Validate job submissions
    - Generate config hashes for reproducibility
    - Create job records in database
    - Submit jobs to the queue
    """
    
    def __init__(self):
        self._status_tracker = get_status_tracker()
    
    def _generate_config_hash(
        self,
        model_name: str,
        model_version: str,
        dataset_name: str,
        dataset_version: str,
        mutation_depth: int,
        attack_types: list[str],
    ) -> str:
        """Generate SHA256 hash of configuration for reproducibility."""
        config = {
            "model_name": model_name,
            "model_version": model_version,
            "dataset_name": dataset_name,
            "dataset_version": dataset_version,
            "mutation_depth": mutation_depth,
            "attack_types": sorted(attack_types),
        }
        config_str = json.dumps(config, sort_keys=True)
        return hashlib.sha256(config_str.encode()).hexdigest()
    
    async def submit_job(
        self,
        request: JobSubmissionRequest,
        submitted_by: Optional[str] = None,
    ) -> EvaluationJob:
        """
        Submit a new evaluation job.
        
        Args:
            request: Job submission request
            submitted_by: API key owner who submitted the job
            
        Returns:
            Created evaluation job
        """
        try:
            # Generate config hash
            config_hash = self._generate_config_hash(
                model_name=request.model_name,
                model_version=request.model_version,
                dataset_name=request.dataset_name,
                dataset_version=request.dataset_version,
                mutation_depth=request.mutation_depth,
                attack_types=request.attack_types,
            )
            
            # Create job
            job = EvaluationJob(
                job_id=uuid.uuid4(),
                job_type=request.job_type,
                model_name=request.model_name,
                model_version=request.model_version,
                dataset_name=request.dataset_name,
                dataset_version=request.dataset_version,
                config_hash=config_hash,
                priority=request.priority,
                submitted_by=submitted_by,
                status=JobStatus.PENDING,
                progress=0.0,
                total_samples=0,  # Will be updated when job starts
                completed_samples=0,
                failed_samples=0,
                checkpoint_interval=request.checkpoint_interval,
                created_at=datetime.utcnow(),
                metadata={
                    "mutation_depth": request.mutation_depth,
                    "attack_types": request.attack_types,
                    "max_concurrency": request.max_concurrency,
                    "sampling_config": request.sampling_config,
                },
            )
            
            # Create job in database
            await self._status_tracker.create_job(job)
            
            # Update status to queued
            await self._status_tracker.update_job_status(
                job.job_id,
                JobStatus.QUEUED,
            )
            job.status = JobStatus.QUEUED
            job.queued_at = datetime.utcnow()
            
            # Add to in-memory queue (for now, using simple list)
            # In production, this would use Redis/RQ/Celery
            _job_queue.append(job)
            
            logger.info(
                "Job submitted",
                job_id=str(job.job_id),
                job_type=job.job_type,
                model=job.model_name,
                dataset=job.dataset_name,
                priority=job.priority,
            )
            
            return job
            
        except Exception as e:
            logger.error(
                "Failed to submit job",
                error=str(e),
            )
            raise
    
    async def submit_benchmark_job(
        self,
        model_name: str,
        model_version: str,
        dataset_version: str,
        submitted_by: Optional[str] = None,
        priority: JobPriority = JobPriority.NORMAL,
        mutation_depth: int = 2,
        attack_types: Optional[list[str]] = None,
        max_concurrency: int = 4,
        checkpoint_interval: int = 10,
    ) -> EvaluationJob:
        """
        Submit a benchmark evaluation job.
        
        Convenience method for submitting standard benchmark jobs.
        
        Args:
            model_name: Name of the model to evaluate
            model_version: Model version
            dataset_version: Dataset version to use
            submitted_by: API key owner
            priority: Job priority
            mutation_depth: Mutation depth
            attack_types: List of attack types
            max_concurrency: Maximum concurrent samples
            checkpoint_interval: Checkpoint interval
            
        Returns:
            Created evaluation job
        """
        request = JobSubmissionRequest(
            job_type=JobType.BENCHMARK,
            model_name=model_name,
            model_version=model_version,
            dataset_name="default",
            dataset_version=dataset_version,
            priority=priority,
            mutation_depth=mutation_depth,
            attack_types=attack_types or ["jailbreak"],
            max_concurrency=max_concurrency,
            checkpoint_interval=checkpoint_interval,
        )
        
        return await self.submit_job(request, submitted_by)
    
    async def cancel_job(
        self,
        job_id: uuid.UUID,
    ) -> bool:
        """
        Cancel a job.
        
        Args:
            job_id: The job ID to cancel
            
        Returns:
            True if cancelled, False if not found or already completed
        """
        try:
            job = self._status_tracker.get_cached_job(job_id)
            
            if job is None:
                logger.warning(
                    "Job not found for cancellation",
                    job_id=str(job_id),
                )
                return False
            
            # Check if job can be cancelled
            if job.status in [JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED]:
                logger.info(
                    "Job already in terminal state",
                    job_id=str(job_id),
                    status=job.status,
                )
                return False
            
            # Update status
            await self._status_tracker.update_job_status(
                job_id,
                JobStatus.CANCELLED,
            )
            
            logger.info(
                "Job cancelled",
                job_id=str(job_id),
            )
            
            return True
            
        except Exception as e:
            logger.error(
                "Failed to cancel job",
                job_id=str(job_id),
                error=str(e),
            )
            return False
    
    def get_queue_size(self) -> int:
        """Get current queue size."""
        return len(_job_queue)
    
    def get_pending_jobs(self) -> list[EvaluationJob]:
        """Get all pending jobs in queue."""
        return [job for job in _job_queue if job.status == JobStatus.QUEUED]


# In-memory job queue (for demonstration)
# In production, this would be replaced with Redis/RQ/Celery
_job_queue: list[EvaluationJob] = []


# Global instance
_producer: Optional[JobProducer] = None


def get_job_producer() -> JobProducer:
    """Get the global job producer instance."""
    global _producer
    if _producer is None:
        _producer = JobProducer()
    return _producer


__all__ = [
    "JobProducer",
    "get_job_producer",
]