"""
async_pipeline.py
-----------------
Async ingestion pipeline for high-throughput episode processing.

Author: HIEROMEM Team
"""

from __future__ import annotations
from typing import Dict, List, Any
import uuid
from celery.result import AsyncResult

from .queue_manager import IngestionQueueManager
from .tasks import process_episode, batch_index_episodes


class AsyncIngestionPipeline:
    """
    Async pipeline for episode ingestion using Celery workers.

    Features:
    - Submit individual or batch episodes
    - Track task status
    - Wait for completion with timeout
    """

    def __init__(
        self,
        redis_url: str = "redis://localhost:6379/1",
        queue_name: str = "hieromem:ingestion",
        debug: bool = False
    ):
        """
        Initialize async pipeline.

        Args:
            redis_url: Redis connection URL
            queue_name: Queue name prefix
            debug: Enable debug logging
        """
        self.queue_manager = IngestionQueueManager(
            redis_url=redis_url,
            queue_name=queue_name,
            debug=debug
        )
        self.debug = debug

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

    def submit_episode(
        self,
        episode_data: Dict[str, Any],
        priority: int = 0,
        use_queue: bool = True
    ) -> str:
        """
        Submit an episode for async processing.

        Args:
            episode_data: Episode data to process
            priority: Processing priority (higher = processed first)
            use_queue: If True, use queue manager; if False, submit directly to Celery

        Returns:
            Task ID
        """
        # Generate task ID if not present
        if 'episode_id' not in episode_data:
            episode_data['episode_id'] = str(uuid.uuid4())

        task_id = episode_data['episode_id']

        if use_queue:
            # Add to queue
            self.queue_manager.enqueue_episode(
                task_id=task_id,
                episode_data=episode_data,
                priority=priority
            )
            self._log(f"Queued episode {task_id}")
        else:
            # Submit directly to Celery
            process_episode.apply_async(
                args=[episode_data],
                task_id=task_id,
                priority=priority
            )
            self._log(f"Submitted episode {task_id} to Celery")

        return task_id

    def submit_batch(
        self,
        episodes: List[Dict[str, Any]],
        batch_size: int = 100
    ) -> List[str]:
        """
        Submit a batch of episodes for processing.

        Args:
            episodes: List of episode data
            batch_size: Size of batches for processing

        Returns:
            List of task IDs
        """
        task_ids = []

        # Process in batches
        for i in range(0, len(episodes), batch_size):
            batch = episodes[i:i + batch_size]

            # Generate batch task ID
            batch_task_id = f"batch_{uuid.uuid4()}"

            # Submit batch task
            batch_index_episodes.apply_async(
                args=[batch],
                task_id=batch_task_id
            )

            task_ids.append(batch_task_id)
            self._log(
                f"Submitted batch {batch_task_id} with {len(batch)} episodes")

        return task_ids

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

        Args:
            task_id: Task identifier

        Returns:
            Task status dictionary
        """
        # Check queue manager first
        task = self.queue_manager.get_task_status(task_id)
        if task:
            return {
                'task_id': task_id,
                'status': task.status,
                'created_at': task.created_at,
                'started_at': task.started_at,
                'completed_at': task.completed_at,
                'error': task.error,
                'retry_count': task.retry_count
            }

        # Check Celery result backend
        result = AsyncResult(task_id)
        return {
            'task_id': task_id,
            'status': result.state,
            'result': result.result if result.ready() else None,
            'error': str(result.info) if result.failed() else None
        }

    def wait_for_completion(
        self,
        task_ids: List[str],
        timeout: int = 300
    ) -> Dict[str, Any]:
        """
        Wait for tasks to complete.

        Args:
            task_ids: List of task IDs to wait for
            timeout: Timeout in seconds

        Returns:
            Dictionary with completion status
        """
        results = {}

        for task_id in task_ids:
            result = AsyncResult(task_id)
            try:
                # Wait for result with timeout
                output = result.get(timeout=timeout)
                results[task_id] = {
                    'status': 'completed',
                    'result': output
                }
            except Exception as e:
                results[task_id] = {
                    'status': 'failed',
                    'error': str(e)
                }

        return results

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


__all__ = ["AsyncIngestionPipeline"]
