"""
queue_manager.py
----------------
Redis-backed queue manager for async episode ingestion.

Provides priority queuing, task tracking, and retry logic for high-throughput
episode processing.

Author: HIEROMEM Team
"""

from __future__ import annotations
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, asdict
import json
import time
import redis
from datetime import datetime
from hieromem_core.circuit_breaker import CircuitBreaker
from hieromem_core.utils.retry import with_retry
from redis.exceptions import ConnectionError, TimeoutError


@dataclass
class IngestionTask:
    """Represents an ingestion task in the queue."""
    task_id: str
    episode_data: Dict[str, Any]
    priority: int
    status: str  # pending, processing, completed, failed
    created_at: str
    started_at: Optional[str] = None
    completed_at: Optional[str] = None
    error: Optional[str] = None
    retry_count: int = 0


class IngestionQueueManager:
    """
    Manages Redis-backed ingestion queue for async episode processing.

    Features:
    - Priority queue using Redis sorted sets
    - Task metadata tracking
    - Dead letter queue for failed tasks
    - Retry logic with exponential backoff
    """

    def __init__(
        self,
        redis_url: str = "redis://localhost:6379/1",
        queue_name: str = "hieromem:ingestion",
        max_retries: int = 3,
        debug: bool = False
    ):
        """
        Initialize queue manager.

        Args:
            redis_url: Redis connection URL
            queue_name: Base name for queue keys
            max_retries: Maximum retry attempts
            debug: Enable debug logging
        """
        self.redis_client = redis.from_url(redis_url, decode_responses=True)
        self.queue_name = queue_name
        self.max_retries = max_retries
        self.debug = debug

        # Queue keys
        self.pending_queue = f"{queue_name}:pending"
        self.processing_set = f"{queue_name}:processing"
        self.completed_set = f"{queue_name}:completed"
        self.failed_queue = f"{queue_name}:failed"
        self.task_metadata = f"{queue_name}:tasks"

        # Circuit breaker for Redis operations
        self.cb = CircuitBreaker("redis_queue", failure_threshold=5, recovery_timeout=30)

        self._log("Initialized ingestion queue manager")

    def _log(self, message: str) -> None:
        """Log debug messages if debug mode is enabled."""
        if self.debug:
            print(f"[IngestionQueueManager] {message}")

    def enqueue_episode(
        self,
        task_id: str,
        episode_data: Dict[str, Any],
        priority: int = 0
    ) -> str:
        """
        Add an episode to the ingestion queue.

        Args:
            task_id: Unique task identifier
            episode_data: Episode data to process
            priority: Priority (higher = processed first)

        Returns:
            Task ID
        """
        task = IngestionTask(
            task_id=task_id,
            episode_data=episode_data,
            priority=priority,
            status="pending",
            created_at=datetime.now().isoformat()
        )

        @with_retry(exceptions=(ConnectionError, TimeoutError))
        def _enqueue_op():
            # Add to pending queue (sorted set by priority)
            self.redis_client.zadd(
                self.pending_queue, {
                    task_id: -priority})  # Negative for descending order

            # Store task metadata
            self.redis_client.hset(
                self.task_metadata,
                task_id,
                json.dumps(
                    asdict(task)))

        self.cb.call(_enqueue_op)
        self._log(f"Enqueued task {task_id} with priority {priority}")
        return task_id

    def dequeue_batch(self, batch_size: int = 10) -> List[IngestionTask]:
        """
        Dequeue a batch of tasks for processing.

        Args:
            batch_size: Number of tasks to dequeue

        Returns:
            List of IngestionTask objects
        """
        @with_retry(exceptions=(ConnectionError, TimeoutError))
        def _dequeue_op():
            # Get highest priority tasks
            task_ids = self.redis_client.zrange(
                self.pending_queue, 0, batch_size - 1)

            if not task_ids:
                return []

            tasks = []
            for task_id in task_ids:
                # Move from pending to processing
                self.redis_client.zrem(self.pending_queue, task_id)
                self.redis_client.sadd(self.processing_set, task_id)

                # Get task metadata
                task_json = self.redis_client.hget(self.task_metadata, task_id)
                if task_json:
                    task_dict = json.loads(task_json)
                    task = IngestionTask(**task_dict)
                    task.status = "processing"
                    task.started_at = datetime.now().isoformat()

                    # Update metadata
                    self.redis_client.hset(
                        self.task_metadata, task_id, json.dumps(
                            asdict(task)))
                    tasks.append(task)
            return tasks

        tasks = self.cb.call(_dequeue_op)
        self._log(f"Dequeued {len(tasks)} tasks")
        return tasks

    def mark_complete(self, task_id: str) -> None:
        """
        Mark a task as completed.

        Args:
            task_id: Task identifier
        """
        @with_retry(exceptions=(ConnectionError, TimeoutError))
        def _mark_complete_op():
            # Move from processing to completed
            self.redis_client.srem(self.processing_set, task_id)
            self.redis_client.sadd(self.completed_set, task_id)

            # Update metadata
            task_json = self.redis_client.hget(self.task_metadata, task_id)
            if task_json:
                task_dict = json.loads(task_json)
                task = IngestionTask(**task_dict)
                task.status = "completed"
                task.completed_at = datetime.now().isoformat()

                self.redis_client.hset(
                    self.task_metadata,
                    task_id,
                    json.dumps(
                        asdict(task)))

        self.cb.call(_mark_complete_op)
        self._log(f"Marked task {task_id} as completed")

    def mark_failed(self, task_id: str, error: str) -> None:
        """
        Mark a task as failed and handle retry logic.

        Args:
            task_id: Task identifier
            error: Error message
        """
        @with_retry(exceptions=(ConnectionError, TimeoutError))
        def _mark_failed_op():
            task_json = self.redis_client.hget(self.task_metadata, task_id)
            if not task_json:
                return

            task_dict = json.loads(task_json)
            task = IngestionTask(**task_dict)
            task.retry_count += 1
            task.error = error

            if task.retry_count < self.max_retries:
                # Retry with exponential backoff
                backoff_priority = task.priority - (2 ** task.retry_count)

                # Move back to pending queue
                self.redis_client.srem(self.processing_set, task_id)
                self.redis_client.zadd(
                    self.pending_queue, {
                        task_id: -backoff_priority})

                task.status = "pending"
                self._log(
                    f"Retrying task {task_id} (attempt {task.retry_count}/{self.max_retries})")
            else:
                # Move to dead letter queue
                self.redis_client.srem(self.processing_set, task_id)
                self.redis_client.zadd(self.failed_queue, {task_id: time.time()})

                task.status = "failed"
                task.completed_at = datetime.now().isoformat()
                self._log(
                    f"Task {task_id} failed after {self.max_retries} retries")

            # Update metadata
            self.redis_client.hset(
                self.task_metadata,
                task_id,
                json.dumps(
                    asdict(task)))

        self.cb.call(_mark_failed_op)

    def get_task_status(self, task_id: str) -> Optional[IngestionTask]:
        """
        Get the status of a task.

        Args:
            task_id: Task identifier

        Returns:
            IngestionTask or None if not found
        """
        @with_retry(exceptions=(ConnectionError, TimeoutError))
        def _get_status_op():
            task_json = self.redis_client.hget(self.task_metadata, task_id)
            if task_json:
                task_dict = json.loads(task_json)
                return IngestionTask(**task_dict)
            return None

        return self.cb.call(_get_status_op)

    def get_queue_stats(self) -> Dict[str, int]:
        """
        Get statistics for the queue.

        Returns:
            Dictionary with queue statistics
        """
        @with_retry(exceptions=(ConnectionError, TimeoutError))
        def _get_stats_op():
            return {
                'pending': self.redis_client.zcard(self.pending_queue),
                'processing': self.redis_client.scard(self.processing_set),
                'completed': self.redis_client.scard(self.completed_set),
                'failed': self.redis_client.zcard(self.failed_queue)
            }
        return self.cb.call(_get_stats_op)

    def clear_completed(self, older_than_seconds: int = 86400) -> int:
        """
        Clear completed tasks older than specified time.

        Args:
            older_than_seconds: Age threshold in seconds (default: 24 hours)

        Returns:
            Number of tasks cleared
        """
        @with_retry(exceptions=(ConnectionError, TimeoutError))
        def _clear_op():
            cutoff_time = time.time() - older_than_seconds

            # Get completed task IDs
            completed_ids = list(self.redis_client.smembers(self.completed_set))

            cleared = 0
            for task_id in completed_ids:
                task_json = self.redis_client.hget(self.task_metadata, task_id)
                if task_json:
                    task_dict = json.loads(task_json)
                    completed_at = datetime.fromisoformat(
                        task_dict.get('completed_at', ''))

                    if completed_at.timestamp() < cutoff_time:
                        self.redis_client.srem(self.completed_set, task_id)
                        self.redis_client.hdel(self.task_metadata, task_id)
                        cleared += 1
            return cleared

        cleared = self.cb.call(_clear_op)
        self._log(f"Cleared {cleared} completed tasks")
        return cleared


__all__ = ["IngestionQueueManager", "IngestionTask"]
