"""
tasks.py
--------
Celery tasks for async episode encoding and indexing.

Author: HIEROMEM Team
"""

from celery import Task
from typing import Dict, Any, Optional
import numpy as np
import os

from .celery_app import app
from hieromem_core.persistence.sharded_faiss_index import ShardedFaissIndex
from hieromem_core.persistence.index_shard_manager import IndexShardManager


class IngestionTask(Task):
    """Base task with shared resources."""

    _shard_manager = None
    _sharded_index = None

    @property
    def shard_manager(self):
        """Lazy-load shard manager."""
        if self._shard_manager is None:
            import os
            base_path = os.getenv("HIEROMEM_SHARD_BASE_PATH", "/data/shards")
            num_shards = int(os.getenv("HIEROMEM_NUM_SHARDS", "4"))

            self._shard_manager = IndexShardManager(
                base_path=base_path,
                num_shards=num_shards,
                debug=True
            )
        return self._shard_manager

    @property
    def sharded_index(self):
        """Lazy-load sharded index."""
        if self._sharded_index is None:
            dimension = int(os.getenv("HIEROMEM_VECTOR_DIMENSION", "384"))

            self._sharded_index = ShardedFaissIndex(
                shard_manager=self.shard_manager,
                dimension=dimension,
                debug=True
            )
        return self._sharded_index


@app.task(base=IngestionTask, max_retries=3, default_retry_delay=60)
def encode_episode(episode_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Encode an episode to generate its embedding vector.

    Args:
        episode_data: Episode data including content

    Returns:
        Episode data with embedding added
    """
    try:
        # Import factory (lazy to avoid loading at module level)
        from hieromem_core.encoders import EncoderFactory

        # Determine media type (default to text)
        media_type = episode_data.get('media_type', 'text')

        # Get appropriate encoder
        encoder = EncoderFactory.get_encoder(media_type)

        # Extract content
        content = episode_data.get('content')
        if content is None:
            raise ValueError("Episode data missing 'content' field")

        # Generate embedding
        embedding = encoder.encode(content)

        # Add to episode data
        if hasattr(embedding, 'tolist'):
            episode_data['embedding'] = embedding.tolist()
        else:
            episode_data['embedding'] = embedding

        return episode_data

    except Exception as exc:
        # Since we are not bound, we cannot retry via self; raise exception
        raise exc


@app.task(bind=True, base=IngestionTask, max_retries=3, default_retry_delay=60)
def index_episode(
    self,
    episode_id: str,
    embedding: list,
    metadata: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Index an episode in the sharded FAISS index.

    Args:
        episode_id: Episode identifier
        embedding: Episode embedding vector
        metadata: Optional metadata

    Returns:
        Result with shard_id and status
    """
    try:
        # Convert embedding to numpy array
        vector = np.array(embedding, dtype=np.float32)

        # Add to sharded index
        shard_id = self.sharded_index.add(episode_id, vector, metadata)

        # Save the shard
        self.sharded_index.save_all_shards()

        return {
            'episode_id': episode_id,
            'shard_id': shard_id,
            'status': 'indexed'
        }

    except Exception as exc:
        # Retry on failure
        raise self.retry(exc=exc)


@app.task(bind=True, base=IngestionTask, max_retries=3, default_retry_delay=60)
def process_episode(self, episode_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Complete pipeline: encode and index an episode.

    Args:
        episode_data: Episode data

    Returns:
        Result with indexing status
    """
    try:
        # Step 1: Encode
        encoded_data = encode_episode(episode_data)

        # Step 2: Index
        result = index_episode(
            episode_id=encoded_data['episode_id'],
            embedding=encoded_data['embedding'],
            metadata=encoded_data.get('metadata')
        )

        return result

    except Exception as exc:
        raise self.retry(exc=exc)


@app.task(bind=True, base=IngestionTask)
def batch_index_episodes(self, episode_batch: list) -> Dict[str, Any]:
    """
    Process a batch of episodes efficiently.

    Args:
        episode_batch: List of episode data dictionaries

    Returns:
        Batch processing results
    """
    try:
        from sentence_transformers import SentenceTransformer
        import os

        model_name = os.getenv("HIEROMEM_ENCODER_MODEL", "all-MiniLM-L6-v2")
        encoder = SentenceTransformer(model_name)

        # Batch encode
        contents = [ep.get('content', '') for ep in episode_batch]
        embeddings = encoder.encode(
            contents,
            convert_to_numpy=True,
            show_progress_bar=False)

        # Batch index
        episode_ids = [ep['episode_id'] for ep in episode_batch]
        metadata_list = [ep.get('metadata') for ep in episode_batch]

        shard_counts = self.sharded_index.add_batch(
            episode_ids=episode_ids,
            vectors=embeddings,
            metadata_list=metadata_list
        )

        # Save all affected shards
        self.sharded_index.save_all_shards()

        return {
            'processed': len(episode_batch),
            'shard_counts': shard_counts,
            'status': 'completed'
        }

    except Exception as exc:
        raise self.retry(exc=exc)


@app.task
def cleanup_old_tasks():
    """Periodic task to clean up old completed tasks."""
    from .queue_manager import IngestionQueueManager
    import os

    redis_url = os.getenv("HIEROMEM_CELERY_BROKER", "redis://localhost:6379/1")
    queue_manager = IngestionQueueManager(redis_url=redis_url)

    # Clear tasks older than 24 hours
    cleared = queue_manager.clear_completed(older_than_seconds=86400)

    return {'cleared': cleared}


__all__ = [
    "encode_episode",
    "index_episode",
    "process_episode",
    "batch_index_episodes",
    "cleanup_old_tasks"
]
