"""Bulk email processing pipeline."""
import logging
from typing import List, Dict, Any, Optional, Callable
from dataclasses import dataclass, field
from pathlib import Path
import json
from datetime import datetime

from src.email_providers.base import Email, ClassificationResult
from src.classification.adaptive_classifier import AdaptiveClassifier

logger = logging.getLogger(__name__)


@dataclass
class ProcessingCheckpoint:
    """Checkpoint for resumable processing."""
    total_emails: int
    processed_count: int
    completed_emails: List[str] = field(default_factory=list)
    queued_for_llm: List[str] = field(default_factory=list)
    failed_emails: List[Dict[str, Any]] = field(default_factory=list)
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            'total_emails': self.total_emails,
            'processed_count': self.processed_count,
            'completed_emails': self.completed_emails,
            'queued_for_llm': self.queued_for_llm,
            'failed_emails': self.failed_emails,
            'timestamp': self.timestamp,
            'progress': f"{self.processed_count}/{self.total_emails}"
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ProcessingCheckpoint":
        """Create from dictionary."""
        return cls(
            total_emails=data.get('total_emails', 0),
            processed_count=data.get('processed_count', 0),
            completed_emails=data.get('completed_emails', []),
            queued_for_llm=data.get('queued_for_llm', []),
            failed_emails=data.get('failed_emails', []),
            timestamp=data.get('timestamp', datetime.now().isoformat())
        )


class BulkProcessor:
    """Process large batches of emails with checkpointing."""

    def __init__(
        self,
        classifier: AdaptiveClassifier,
        batch_size: int = 100,
        checkpoint_dir: str = "checkpoints",
        checkpoint_interval: int = 1000
    ):
        """Initialize bulk processor."""
        self.classifier = classifier
        self.batch_size = batch_size
        self.checkpoint_dir = checkpoint_dir
        self.checkpoint_interval = checkpoint_interval
        self.checkpoint = None

        Path(self.checkpoint_dir).mkdir(parents=True, exist_ok=True)

    def process(
        self,
        emails: List[Email],
        resume: bool = True,
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> tuple[List[ClassificationResult], List[str]]:
        """
        Process batch of emails with checkpointing.

        Args:
            emails: List of emails to process
            resume: Try to resume from checkpoint if exists
            progress_callback: Function to call with (processed, total)

        Returns:
            (classified_results, llm_queue_ids)
        """
        logger.info(f"Starting bulk processing of {len(emails)} emails")

        # Try to load checkpoint
        if resume:
            self.checkpoint = self._load_checkpoint(len(emails))
            if self.checkpoint:
                logger.info(f"Resumed from checkpoint: {self.checkpoint.processed_count}/{len(emails)} done")
            else:
                self.checkpoint = ProcessingCheckpoint(total_emails=len(emails), processed_count=0)
        else:
            self.checkpoint = ProcessingCheckpoint(total_emails=len(emails), processed_count=0)

        results = []
        llm_queue = []

        # Process in batches
        start_idx = self.checkpoint.processed_count
        for batch_idx in range(start_idx, len(emails), self.batch_size):
            batch_end = min(batch_idx + self.batch_size, len(emails))
            batch = emails[batch_idx:batch_end]

            logger.debug(f"Processing batch {batch_idx}-{batch_end}")

            # Process batch
            for email in batch:
                try:
                    # Skip if already processed
                    if email.id in self.checkpoint.completed_emails:
                        logger.debug(f"Skipping already processed: {email.id}")
                        continue

                    # Classify
                    result = self.classifier.classify(email)
                    results.append(result)

                    # Track for LLM if needed
                    if result.needs_review:
                        llm_queue.append(email.id)
                        self.checkpoint.queued_for_llm.append(email.id)

                    self.checkpoint.completed_emails.append(email.id)
                    self.checkpoint.processed_count += 1

                except Exception as e:
                    logger.error(f"Error processing {email.id}: {e}")
                    self.checkpoint.failed_emails.append({
                        'email_id': email.id,
                        'error': str(e),
                        'timestamp': datetime.now().isoformat()
                    })

            # Checkpoint periodically
            if (batch_end - start_idx) % self.checkpoint_interval == 0:
                self._save_checkpoint()
                logger.info(f"Checkpoint saved: {self.checkpoint.processed_count}/{len(emails)}")

            # Progress callback
            if progress_callback:
                progress_callback(self.checkpoint.processed_count, len(emails))

        # Final checkpoint
        self._save_checkpoint()
        logger.info(f"Processing complete: {len(results)} classified, {len(llm_queue)} need LLM review")

        return results, llm_queue

    def _save_checkpoint(self) -> None:
        """Save current checkpoint."""
        if not self.checkpoint:
            return

        checkpoint_file = Path(self.checkpoint_dir) / f"checkpoint_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"

        try:
            with open(checkpoint_file, 'w') as f:
                json.dump(self.checkpoint.to_dict(), f, indent=2)
            logger.debug(f"Checkpoint saved to {checkpoint_file}")
        except Exception as e:
            logger.error(f"Error saving checkpoint: {e}")

    def _load_checkpoint(self, expected_total: int) -> Optional[ProcessingCheckpoint]:
        """Load latest checkpoint if it matches."""
        checkpoint_dir = Path(self.checkpoint_dir)

        if not checkpoint_dir.exists():
            return None

        # Find latest checkpoint
        checkpoints = sorted(checkpoint_dir.glob('checkpoint_*.json'), reverse=True)

        if not checkpoints:
            return None

        try:
            latest = checkpoints[0]
            with open(latest, 'r') as f:
                data = json.load(f)

            # Verify checkpoint matches this job
            if data.get('total_emails') == expected_total:
                logger.info(f"Loaded checkpoint from {latest.name}")
                return ProcessingCheckpoint.from_dict(data)
            else:
                logger.warning(f"Checkpoint total mismatch: {data.get('total_emails')} vs {expected_total}")
                return None

        except Exception as e:
            logger.error(f"Error loading checkpoint: {e}")
            return None

    def get_checkpoint_status(self) -> Dict[str, Any]:
        """Get current checkpoint status."""
        if not self.checkpoint:
            return {'status': 'no_checkpoint'}

        return {
            'status': 'active',
            'processed': self.checkpoint.processed_count,
            'total': self.checkpoint.total_emails,
            'progress_percent': (self.checkpoint.processed_count / self.checkpoint.total_emails * 100) if self.checkpoint.total_emails > 0 else 0,
            'queued_for_llm': len(self.checkpoint.queued_for_llm),
            'failed': len(self.checkpoint.failed_emails),
            'timestamp': self.checkpoint.timestamp
        }
