"""Job batching and bulk operations for Oclawma.

This module provides batch job management with:
- Batch grouping of related jobs
- Shared configuration across batch jobs
- Bulk operations (retry, cancel, delete)
- Partial failure handling
- Progress tracking
"""

from __future__ import annotations

import json
import logging
import threading
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Callable

from pydantic import Field

from oclawma.queue.models import Job, JobPriority, JobStatus
from oclawma.queue.queue import JobQueue

logger = logging.getLogger(__name__)


class BatchStatus(str, Enum):
    """Batch status enumeration."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    PARTIAL = "partial"  # Some jobs succeeded, some failed
    FAILED = "failed"  # All jobs failed
    CANCELLED = "cancelled"


class BatchJob(Job):
    """A job that belongs to a batch.

    Extends the base Job model with batch-specific fields.

    Attributes:
        batch_id: Unique identifier for the batch this job belongs to
        batch_config: Shared configuration for all jobs in the batch
        batch_order: Execution order within the batch (for sequential batches)
        batch_result: Result data from job execution
    """

    # Inherit from Job but add batch-specific fields
    batch_id: str | None = Field(default=None, description="Batch identifier")
    batch_config: dict[str, Any] = Field(
        default_factory=dict, description="Shared batch configuration"
    )
    batch_order: int = Field(default=0, description="Execution order within batch")
    batch_result: dict[str, Any] | None = Field(default=None, description="Job result data")

    def to_db_dict(self) -> dict[str, Any]:
        """Convert to dictionary with batch fields embedded in payload."""
        # Get base job dict from parent class

        def fmt_dt(dt: datetime | None) -> str | None:
            if dt is None:
                return None
            return dt.strftime("%Y-%m-%d %H:%M:%S")

        # Store batch fields in payload for persistence
        payload = dict(self.payload)
        payload["__batch__"] = {
            "batch_id": self.batch_id,
            "batch_config": self.batch_config,
            "batch_order": self.batch_order,
            "batch_result": self.batch_result,
        }

        return {
            "id": self.id,
            "payload": json.dumps(payload),
            "status": self.status.value,
            "created_at": fmt_dt(self.created_at),
            "updated_at": fmt_dt(self.updated_at),
            "scheduled_at": fmt_dt(self.scheduled_at),
            "started_at": fmt_dt(self.started_at),
            "completed_at": fmt_dt(self.completed_at),
            "retry_count": self.retry_count,
            "max_retries": self.max_retries,
            "error": self.error,
            "priority": (
                self.priority.value if isinstance(self.priority, JobPriority) else self.priority
            ),
            "depends_on": json.dumps(self.depends_on) if self.depends_on else None,
            "job_type": self.job_type,
            "encrypted": self.encrypted,
            "encryption_metadata": (
                json.dumps(self.encryption_metadata) if self.encryption_metadata else None
            ),
            "cron_expression": self.cron_expression,
            "timezone": self.timezone,
            "next_run_at": fmt_dt(self.next_run_at),
        }

    @classmethod
    def from_job(cls, job: Job) -> BatchJob:
        """Create a BatchJob from a regular Job, extracting batch data from payload."""
        # Extract batch fields from payload if present
        batch_id = None
        batch_config: dict[str, Any] = {}
        batch_order = 0
        batch_result = None

        payload = dict(job.payload)
        if "__batch__" in payload:
            batch_data = payload.pop("__batch__")
            batch_id = batch_data.get("batch_id")
            batch_config = batch_data.get("batch_config", {})
            batch_order = batch_data.get("batch_order", 0)
            batch_result = batch_data.get("batch_result")

        return cls(
            id=job.id,
            payload=payload,
            status=job.status,
            created_at=job.created_at,
            updated_at=job.updated_at,
            scheduled_at=job.scheduled_at,
            started_at=job.started_at,
            completed_at=job.completed_at,
            retry_count=job.retry_count,
            max_retries=job.max_retries,
            error=job.error,
            priority=job.priority,
            depends_on=job.depends_on,
            job_type=job.job_type,
            encrypted=job.encrypted,
            encryption_metadata=job.encryption_metadata,
            cron_expression=job.cron_expression,
            timezone=job.timezone,
            next_run_at=job.next_run_at,
            batch_id=batch_id,
            batch_config=batch_config,
            batch_order=batch_order,
            batch_result=batch_result,
        )


@dataclass
class BatchProgress:
    """Progress information for a batch.

    Attributes:
        batch_id: Unique batch identifier
        total: Total number of jobs in the batch
        pending: Number of pending jobs
        running: Number of running jobs
        completed: Number of completed jobs
        failed: Number of failed jobs
        cancelled: Number of cancelled jobs
        status: Overall batch status
        percent_complete: Percentage of completion (0-100)
    """

    batch_id: str
    total: int = 0
    pending: int = 0
    running: int = 0
    completed: int = 0
    failed: int = 0
    cancelled: int = 0
    status: BatchStatus = BatchStatus.PENDING
    percent_complete: float = 0.0
    errors: list[str] = field(default_factory=list)

    def __post_init__(self):
        """Calculate derived fields."""
        if self.total > 0:
            self.percent_complete = (
                (self.completed + self.failed + self.cancelled) / self.total
            ) * 100

            # Determine overall status
            if self.cancelled == self.total:
                self.status = BatchStatus.CANCELLED
            elif self.failed == self.total:
                self.status = BatchStatus.FAILED
            elif self.completed == self.total:
                self.status = BatchStatus.COMPLETED
            elif self.failed > 0 or self.cancelled > 0:
                self.status = BatchStatus.PARTIAL
            elif self.running > 0:
                self.status = BatchStatus.RUNNING
            else:
                self.status = BatchStatus.PENDING
        else:
            self.status = BatchStatus.PENDING


@dataclass
class BatchResult:
    """Result of a batch operation.

    Attributes:
        batch_id: Unique batch identifier
        success: Whether the operation was successful
        processed_count: Number of jobs processed
        failed_count: Number of jobs that failed
        failed_job_ids: List of job IDs that failed
        errors: List of error messages
    """

    batch_id: str
    success: bool = True
    processed_count: int = 0
    failed_count: int = 0
    failed_job_ids: list[int] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)


class BatchError(Exception):
    """Raised when there's an error with batch operations."""

    pass


class BatchNotFoundError(BatchError):
    """Raised when a batch is not found."""

    pass


class BatchManager:
    """Manager for batch job operations.

    Provides functionality to:
    - Create and manage batches of related jobs
    - Submit multiple jobs as a batch
    - Track batch progress
    - Perform bulk operations on batch jobs
    - Handle partial failures

    Example:
        >>> manager = BatchManager(queue)
        >>> batch_id = manager.create_batch("email-campaign")
        >>> jobs = [
        ...     {"to": "user1@example.com", "subject": "Hello"},
        ...     {"to": "user2@example.com", "subject": "Hello"},
        ... ]
        >>> manager.submit_jobs(batch_id, jobs, shared_config={"template": "welcome"})
        >>> progress = manager.get_progress(batch_id)
        >>> print(f"{progress.percent_complete:.1f}% complete")
    """

    def __init__(self, queue: JobQueue) -> None:
        """Initialize the batch manager.

        Args:
            queue: The JobQueue instance to use for job operations
        """
        self.queue = queue
        self._lock = threading.RLock()

    def create_batch(
        self,
        name: str | None = None,
        description: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """Create a new batch.

        Args:
            name: Optional batch name for identification
            description: Optional batch description
            metadata: Optional metadata dictionary

        Returns:
            Unique batch identifier (UUID)
        """
        batch_id = str(uuid.uuid4())

        # Store batch metadata in a special tracking job
        tracking_payload = {
            "__batch_tracking__": True,
            "batch_id": batch_id,
            "name": name or f"batch-{batch_id[:8]}",
            "description": description,
            "metadata": metadata or {},
            "created_at": datetime.utcnow().isoformat(),
        }

        with self._lock:
            self.queue.enqueue(
                payload=tracking_payload,
                priority=JobPriority.NORMAL,
                job_type="batch_tracking",
            )

        logger.info(f"Created batch {batch_id} with name '{name}'")
        return batch_id

    def submit_jobs(
        self,
        batch_id: str,
        payloads: list[dict[str, Any]],
        shared_config: dict[str, Any] | None = None,
        priority: JobPriority | int = JobPriority.NORMAL,
        max_retries: int | None = None,
        sequential: bool = False,
    ) -> list[BatchJob]:
        """Submit multiple jobs as a batch.

        Args:
            batch_id: The batch identifier to associate jobs with
            payloads: List of job payloads to submit
            shared_config: Configuration shared across all jobs in the batch
            priority: Priority level for all jobs in the batch
            max_retries: Maximum retries for each job
            sequential: If True, jobs will have dependencies to run sequentially

        Returns:
            List of created BatchJob instances

        Raises:
            BatchError: If batch_id is invalid or submission fails
        """
        if not batch_id:
            raise BatchError("batch_id is required")

        if not payloads:
            return []

        shared_config = shared_config or {}
        jobs: list[BatchJob] = []
        previous_job_id: int | None = None

        with self._lock:
            for idx, payload in enumerate(payloads):
                # Create depends_on for sequential execution
                depends_on = [previous_job_id] if sequential and previous_job_id else None

                # Create the base job
                base_job = self.queue.enqueue(
                    payload=payload,
                    priority=priority,
                    max_retries=max_retries,
                    depends_on=depends_on,
                )

                # Convert to batch job by updating payload with batch data
                batch_job = BatchJob(
                    id=base_job.id,
                    payload=payload,
                    status=base_job.status,
                    created_at=base_job.created_at,
                    updated_at=base_job.updated_at,
                    scheduled_at=base_job.scheduled_at,
                    started_at=base_job.started_at,
                    completed_at=base_job.completed_at,
                    retry_count=base_job.retry_count,
                    max_retries=base_job.max_retries,
                    error=base_job.error,
                    priority=base_job.priority,
                    depends_on=base_job.depends_on,
                    job_type=base_job.job_type,
                    encrypted=base_job.encrypted,
                    encryption_metadata=base_job.encryption_metadata,
                    batch_id=batch_id,
                    batch_config=shared_config,
                    batch_order=idx,
                )

                # Update the job with batch info in payload
                self._update_job_with_batch_data(batch_job)
                jobs.append(batch_job)

                if sequential:
                    previous_job_id = base_job.id

        logger.info(f"Submitted {len(jobs)} jobs to batch {batch_id}")
        return jobs

    def _update_job_with_batch_data(self, batch_job: BatchJob) -> None:
        """Update a job's payload with batch data in the database."""
        try:
            job = self.queue.get_job(batch_job.id, decrypt=False)
            # Embed batch data in payload
            job.payload["__batch__"] = {
                "batch_id": batch_job.batch_id,
                "batch_config": batch_job.batch_config,
                "batch_order": batch_job.batch_order,
                "batch_result": batch_job.batch_result,
            }
            self.queue.store.update(job)
        except Exception as e:
            logger.warning(f"Failed to update job {batch_job.id} with batch data: {e}")

    def get_progress(self, batch_id: str) -> BatchProgress:
        """Get the progress of a batch.

        Args:
            batch_id: The batch identifier

        Returns:
            BatchProgress with current status

        Raises:
            BatchNotFoundError: If batch doesn't exist
        """
        jobs = self._get_batch_jobs(batch_id)

        if not jobs:
            raise BatchNotFoundError(f"Batch {batch_id} not found or has no jobs")

        total = len(jobs)
        pending = sum(1 for j in jobs if j.status == JobStatus.PENDING)
        running = sum(1 for j in jobs if j.status == JobStatus.RUNNING)
        completed = sum(1 for j in jobs if j.status == JobStatus.COMPLETED)
        failed = sum(1 for j in jobs if j.status == JobStatus.FAILED)
        # Note: Cancelled jobs are deleted from the queue, so we track them separately
        # For now, cancelled is always 0 since cancelled jobs don't exist in the store
        cancelled = 0

        errors = [j.error for j in jobs if j.error]

        return BatchProgress(
            batch_id=batch_id,
            total=total,
            pending=pending,
            running=running,
            completed=completed,
            failed=failed,
            cancelled=cancelled,
            errors=errors,
        )

    def _get_batch_jobs(self, batch_id: str) -> list[BatchJob]:
        """Get all jobs belonging to a batch.

        Args:
            batch_id: The batch identifier

        Returns:
            List of BatchJob instances
        """
        jobs = []
        # Get all jobs and filter by batch_id
        all_jobs = self.queue.list_jobs(decrypt=False)

        for job in all_jobs:
            if "__batch__" in job.payload:
                batch_data = job.payload["__batch__"]
                if batch_data.get("batch_id") == batch_id:
                    jobs.append(BatchJob.from_job(job))

        return jobs

    def list_batches(self) -> list[dict[str, Any]]:
        """List all batches with their progress.

        Returns:
            List of batch summaries with progress information
        """
        # Find all batch tracking jobs
        all_jobs = self.queue.list_jobs(decrypt=False)
        batches: dict[str, dict[str, Any]] = {}

        # First, find batch tracking jobs
        for job in all_jobs:
            if job.payload.get("__batch_tracking__"):
                batch_id = job.payload.get("batch_id")
                if batch_id:
                    batches[batch_id] = {
                        "batch_id": batch_id,
                        "name": job.payload.get("name", f"batch-{batch_id[:8]}"),
                        "description": job.payload.get("description"),
                        "created_at": job.payload.get("created_at"),
                        "progress": None,  # Will be populated later
                    }

        # Get jobs for each batch and calculate progress
        for job in all_jobs:
            if "__batch__" in job.payload:
                batch_data = job.payload["__batch__"]
                batch_id = batch_data.get("batch_id")
                if batch_id and batch_id in batches:
                    if "job_count" not in batches[batch_id]:
                        batches[batch_id]["job_count"] = 0
                        batches[batch_id]["completed"] = 0
                        batches[batch_id]["failed"] = 0
                    batches[batch_id]["job_count"] += 1
                    if job.status == JobStatus.COMPLETED:
                        batches[batch_id]["completed"] += 1
                    elif job.status == JobStatus.FAILED:
                        batches[batch_id]["failed"] += 1

        # Calculate progress percentages
        for _batch_id, batch_info in batches.items():
            total = batch_info.get("job_count", 0)
            if total > 0:
                completed = batch_info.get("completed", 0)
                failed = batch_info.get("failed", 0)
                batch_info["percent_complete"] = ((completed + failed) / total) * 100
            else:
                batch_info["percent_complete"] = 0.0

        return list(batches.values())

    # -------------------------------------------------------------------------
    # Bulk Operations
    # -------------------------------------------------------------------------

    def bulk_retry(
        self,
        batch_id: str,
        failed_only: bool = True,
        on_error: Callable[[int, Exception], None] | None = None,
    ) -> BatchResult:
        """Retry failed jobs in a batch.

        Args:
            batch_id: The batch identifier
            failed_only: If True, only retry failed jobs. If False, retry all.
            on_error: Optional callback for per-job errors

        Returns:
            BatchResult with retry results
        """
        jobs = self._get_batch_jobs(batch_id)

        if not jobs:
            raise BatchNotFoundError(f"Batch {batch_id} not found")

        result = BatchResult(batch_id=batch_id)

        for job in jobs:
            should_retry = True
            if failed_only and job.status != JobStatus.FAILED:
                should_retry = False

            if should_retry and job.is_retryable():
                try:
                    self.queue.retry(job.id)
                    result.processed_count += 1
                except Exception as e:
                    result.failed_count += 1
                    result.failed_job_ids.append(job.id)
                    result.errors.append(f"Job {job.id}: {e}")
                    if on_error:
                        on_error(job.id, e)

        result.success = result.failed_count == 0
        logger.info(
            f"Bulk retry for batch {batch_id}: "
            f"{result.processed_count} retried, {result.failed_count} failed"
        )
        return result

    def bulk_cancel(
        self,
        batch_id: str,
        pending_only: bool = False,
        on_error: Callable[[int, Exception], None] | None = None,
    ) -> BatchResult:
        """Cancel jobs in a batch.

        Args:
            batch_id: The batch identifier
            pending_only: If True, only cancel pending jobs.
                         If False, cancel pending, running, and paused jobs.
            on_error: Optional callback for per-job errors

        Returns:
            BatchResult with cancellation results
        """
        jobs = self._get_batch_jobs(batch_id)

        if not jobs:
            raise BatchNotFoundError(f"Batch {batch_id} not found")

        result = BatchResult(batch_id=batch_id)
        cancellable_statuses = {JobStatus.PENDING}

        if not pending_only:
            cancellable_statuses.update({JobStatus.RUNNING, JobStatus.PAUSED})

        for job in jobs:
            if job.status in cancellable_statuses:
                try:
                    if self.queue.cancel(job.id):
                        result.processed_count += 1
                    else:
                        result.failed_count += 1
                        result.failed_job_ids.append(job.id)
                except Exception as e:
                    result.failed_count += 1
                    result.failed_job_ids.append(job.id)
                    result.errors.append(f"Job {job.id}: {e}")
                    if on_error:
                        on_error(job.id, e)

        result.success = result.failed_count == 0
        logger.info(
            f"Bulk cancel for batch {batch_id}: "
            f"{result.processed_count} cancelled, {result.failed_count} failed"
        )
        return result

    def bulk_delete(
        self,
        batch_id: str,
        completed_only: bool = False,
        on_error: Callable[[int, Exception], None] | None = None,
    ) -> BatchResult:
        """Delete jobs in a batch.

        Args:
            batch_id: The batch identifier
            completed_only: If True, only delete completed jobs.
                          If False, delete all jobs in the batch.
            on_error: Optional callback for per-job errors

        Returns:
            BatchResult with deletion results
        """
        jobs = self._get_batch_jobs(batch_id)

        if not jobs:
            raise BatchNotFoundError(f"Batch {batch_id} not found")

        result = BatchResult(batch_id=batch_id)

        for job in jobs:
            if completed_only and job.status != JobStatus.COMPLETED:
                continue

            try:
                if self.queue.cancel(job.id):  # cancel() also deletes
                    result.processed_count += 1
                else:
                    # Try direct store deletion if cancel failed
                    self.queue.store.delete(job.id)
                    result.processed_count += 1
            except Exception as e:
                result.failed_count += 1
                result.failed_job_ids.append(job.id)
                result.errors.append(f"Job {job.id}: {e}")
                if on_error:
                    on_error(job.id, e)

        result.success = result.failed_count == 0
        logger.info(
            f"Bulk delete for batch {batch_id}: "
            f"{result.processed_count} deleted, {result.failed_count} failed"
        )
        return result

    def get_batch_stats(self, batch_id: str) -> dict[str, Any]:
        """Get detailed statistics for a batch.

        Args:
            batch_id: The batch identifier

        Returns:
            Dictionary with batch statistics
        """
        progress = self.get_progress(batch_id)

        # Calculate additional stats
        jobs = self._get_batch_jobs(batch_id)
        avg_retries = sum(j.retry_count for j in jobs) / len(jobs) if jobs else 0
        total_processing_time = 0.0
        processing_count = 0

        for job in jobs:
            if job.started_at and job.completed_at:
                duration = (job.completed_at - job.started_at).total_seconds()
                total_processing_time += duration
                processing_count += 1

        return {
            "batch_id": batch_id,
            "total_jobs": progress.total,
            "pending": progress.pending,
            "running": progress.running,
            "completed": progress.completed,
            "failed": progress.failed,
            "cancelled": progress.cancelled,
            "status": progress.status.value,
            "percent_complete": progress.percent_complete,
            "average_retries": avg_retries,
            "average_processing_time": (
                total_processing_time / processing_count if processing_count > 0 else 0
            ),
            "total_processing_time": total_processing_time,
            "errors": progress.errors,
        }

    def pause_batch(self, batch_id: str) -> BatchResult:
        """Pause all pending jobs in a batch.

        Args:
            batch_id: The batch identifier

        Returns:
            BatchResult with pause results
        """
        # Note: The current Job model doesn't have a PAUSED status for pending jobs
        # This would require setting a very far future scheduled_at
        # For now, we'll just return info about what would be paused
        jobs = self._get_batch_jobs(batch_id)

        if not jobs:
            raise BatchNotFoundError(f"Batch {batch_id} not found")

        result = BatchResult(batch_id=batch_id)

        for job in jobs:
            if job.status == JobStatus.PENDING:
                # Reschedule to far future (effectively pausing)
                try:
                    store_job = self.queue.store.get(job.id)
                    from datetime import timedelta

                    store_job.scheduled_at = datetime.utcnow() + timedelta(days=365)
                    self.queue.store.update(store_job)
                    result.processed_count += 1
                except Exception as e:
                    result.failed_count += 1
                    result.errors.append(f"Job {job.id}: {e}")

        return result

    def resume_batch(self, batch_id: str) -> BatchResult:
        """Resume all paused jobs in a batch.

        Args:
            batch_id: The batch identifier

        Returns:
            BatchResult with resume results
        """
        jobs = self._get_batch_jobs(batch_id)

        if not jobs:
            raise BatchNotFoundError(f"Batch {batch_id} not found")

        result = BatchResult(batch_id=batch_id)

        for job in jobs:
            # Find jobs that are scheduled far in the future
            # (our way of marking them as paused)
            if job.scheduled_at:
                days_until = (job.scheduled_at - datetime.utcnow()).days
                if days_until > 300:  # Likely paused by us
                    try:
                        store_job = self.queue.store.get(job.id)
                        store_job.scheduled_at = None  # Make it immediately available
                        self.queue.store.update(store_job)
                        result.processed_count += 1
                    except Exception as e:
                        result.failed_count += 1
                        result.errors.append(f"Job {job.id}: {e}")

        return result

    def __enter__(self) -> BatchManager:
        """Context manager entry."""
        return self

    def __exit__(self, *args: Any) -> None:
        """Context manager exit."""
        pass
