"""Queue management for LLM batch processing."""
import logging
import json
from typing import List, Dict, Any, Optional
from pathlib import Path
from datetime import datetime
from collections import deque

logger = logging.getLogger(__name__)


class LLMQueue:
    """
    Queue for emails awaiting LLM review.

    Manages:
    - Batching emails for efficient LLM processing
    - Persistence to disk
    - Processing status tracking
    """

    def __init__(self, batch_size: int = 50, persist_dir: Optional[str] = None):
        """Initialize queue."""
        self.batch_size = batch_size
        self.persist_dir = Path(persist_dir) if persist_dir else None
        self.queue: deque = deque()
        self.processing: Dict[str, Any] = {}  # Currently processing batch
        self.completed: List[str] = []  # Completed email IDs
        self.failed: List[Dict[str, Any]] = []  # Failed items

        if self.persist_dir:
            self.persist_dir.mkdir(parents=True, exist_ok=True)

    def add(self, email_id: str, email_data: Dict[str, Any]) -> None:
        """Add email to queue."""
        self.queue.append({
            'email_id': email_id,
            'data': email_data,
            'added_at': datetime.now().isoformat(),
            'retries': 0
        })

        logger.debug(f"Added to LLM queue: {email_id} (queue size: {len(self.queue)})")

    def add_batch(self, items: List[tuple[str, Dict[str, Any]]]) -> None:
        """Add multiple items to queue."""
        for email_id, email_data in items:
            self.add(email_id, email_data)

        logger.info(f"Added {len(items)} items to LLM queue (total: {len(self.queue)})")

    def get_batch(self) -> Optional[List[Dict[str, Any]]]:
        """Get next batch for processing."""
        if not self.queue:
            return None

        batch = []
        for _ in range(min(self.batch_size, len(self.queue))):
            batch.append(self.queue.popleft())

        self.processing = {
            'batch_id': f"batch_{datetime.now().isoformat()}",
            'items': batch,
            'started_at': datetime.now().isoformat()
        }

        logger.info(f"Got batch for processing: {len(batch)} items (queue remaining: {len(self.queue)})")
        return batch

    def mark_completed(self, email_ids: List[str]) -> None:
        """Mark items as successfully processed."""
        for email_id in email_ids:
            if email_id not in self.completed:
                self.completed.append(email_id)

        logger.debug(f"Marked {len(email_ids)} items completed (total completed: {len(self.completed)})")

    def mark_failed(self, email_id: str, error: str, retry: bool = True) -> bool:
        """
        Mark item as failed.

        Returns:
            True if should retry, False if should give up
        """
        # Find in processing batch
        for item in self.processing.get('items', []):
            if item['email_id'] == email_id:
                item['retries'] += 1

                if item['retries'] < 3 and retry:
                    # Requeue for retry
                    self.queue.append(item)
                    logger.warning(f"Requeuing {email_id} (retry {item['retries']}/3)")
                    return True
                else:
                    # Give up
                    self.failed.append({
                        'email_id': email_id,
                        'error': error,
                        'retries': item['retries'],
                        'failed_at': datetime.now().isoformat()
                    })
                    logger.error(f"Failed to process {email_id} after {item['retries']} retries: {error}")
                    return False

        logger.warning(f"Failed item {email_id} not found in processing batch")
        return False

    def persist(self, filepath: str) -> bool:
        """Save queue state to disk."""
        try:
            state = {
                'timestamp': datetime.now().isoformat(),
                'queue_size': len(self.queue),
                'queue': list(self.queue),
                'processing': self.processing,
                'completed': self.completed,
                'failed': self.failed
            }

            with open(filepath, 'w') as f:
                json.dump(state, f, indent=2)

            logger.debug(f"Queue state persisted to {filepath}")
            return True

        except Exception as e:
            logger.error(f"Error persisting queue: {e}")
            return False

    def restore(self, filepath: str) -> bool:
        """Restore queue state from disk."""
        try:
            with open(filepath, 'r') as f:
                state = json.load(f)

            self.queue = deque(state.get('queue', []))
            self.processing = state.get('processing', {})
            self.completed = state.get('completed', [])
            self.failed = state.get('failed', [])

            logger.info(f"Queue state restored from {filepath}")
            logger.info(f"Queue size: {len(self.queue)}, Completed: {len(self.completed)}, Failed: {len(self.failed)}")
            return True

        except Exception as e:
            logger.error(f"Error restoring queue: {e}")
            return False

    def get_stats(self) -> Dict[str, Any]:
        """Get queue statistics."""
        return {
            'queue_size': len(self.queue),
            'processing': len(self.processing.get('items', [])),
            'completed': len(self.completed),
            'failed': len(self.failed),
            'batch_size': self.batch_size,
            'total_processed': len(self.completed) + len(self.failed)
        }

    def is_empty(self) -> bool:
        """Check if queue is empty."""
        return len(self.queue) == 0 and not self.processing.get('items')

    def clear(self) -> None:
        """Clear all queues."""
        self.queue.clear()
        self.processing = {}
        self.completed = []
        self.failed = []
        logger.info("Queue cleared")
