"""
Training Manager for CyborgDB Service

This module manages automatic index training/retraining based on vector count thresholds.
It uses a queue-based approach with on-demand worker threads to handle training requests.
"""

import threading
import queue
import logging
import os
from typing import Optional, Dict, Any
from dataclasses import dataclass

from cyborgdb_service.core.config import settings
from cyborgdb_service.core.security import hex_to_bytes
from cyborgdb_service.db.client import clear_index_cache
from cyborgdb_core import Client, DBConfig, GPUConfig

logger = logging.getLogger(__name__)


@dataclass
class TrainingJob:
    """Represents a training job to be processed."""

    index_name: str
    index_key_hex: str
    n_lists: Optional[int] = None
    batch_size: Optional[int] = None
    max_iters: Optional[int] = None
    tolerance: Optional[float] = None
    max_memory: Optional[int] = None
    # Set by train_sync(wait=True) so the caller can block until completion
    _completion_event: Optional[threading.Event] = None
    _completion_error: Optional[str] = None


class TrainingManager:
    """Manages index training operations with queue-based processing."""

    # Default retraining threshold - trigger when num_vectors > n_lists * threshold
    DEFAULT_RETRAIN_THRESHOLD = 10000

    def __init__(self):
        """Initialize the training manager."""
        self._lock = threading.Lock()
        self._queue: queue.Queue[TrainingJob] = queue.Queue()
        self._currently_training: Optional[str] = None
        self._worker_running = False
        self._training_client: Optional[Client] = None
        self._retrain_threshold = int(
            os.getenv("RETRAIN_THRESHOLD", self.DEFAULT_RETRAIN_THRESHOLD)
        )

        # Initialize the dedicated training client
        self._initialize_training_client()

        logger.info(
            f"TrainingManager initialized with retrain threshold: {self._retrain_threshold}"
        )

    def _initialize_training_client(self):
        """Initialize a dedicated client for training operations with resource limits."""
        try:
            # Calculate resource limits for training client
            # Use a small portion of available resources to not degrade server performance
            training_cpu_threads = (
                max(1, settings.CPU_THREADS // 4) if settings.CPU_THREADS > 0 else 2
            )

            index_location = DBConfig(
                location=settings.INDEX_LOCATION,
                table_name=settings.INDEX_TABLE_NAME,
                connection_string=settings.INDEX_CONNECTION_STRING,
            )
            config_location = DBConfig(
                location=settings.CONFIG_LOCATION,
                table_name=settings.CONFIG_TABLE_NAME,
                connection_string=settings.CONFIG_CONNECTION_STRING,
            )
            items_location = DBConfig(
                location=settings.ITEMS_LOCATION,
                table_name=settings.ITEMS_TABLE_NAME,
                connection_string=settings.ITEMS_CONNECTION_STRING,
            )

            gpu_config = GPUConfig(upsert=settings.GPU_UPSERT, train=settings.GPU_TRAIN)

            self._training_client = Client(
                api_key=settings.CYBORGDB_API_KEY,
                index_location=index_location,
                config_location=config_location,
                items_location=items_location,
                cpu_threads=training_cpu_threads,
                gpu_config=gpu_config,
            )

            logger.info(
                f"Training client initialized with {training_cpu_threads} CPU threads - "
                f"GPU enabled: {self._training_client.is_gpu_enabled()}, "
                f"GPU config: {self._training_client.get_gpu_config()}"
            )

        except Exception as e:
            logger.error(f"Failed to initialize training client: {str(e)}")
            raise

    def should_retrain(
        self,
        index_name: str,
        num_vectors: int,
        n_lists: int,
        is_trained: bool,
        index_type: str,
    ) -> bool:
        """
        Check if an index should be retrained based on the heuristic and index type.

        Args:
            index_name: Name of the index
            num_vectors: Current number of vectors in the index
            n_lists: Number of lists in the index configuration
            is_trained: Whether the index is currently trained
            index_type: Type of the index (ivf, ivfpq, ivfflat, ivfsq)

        Returns:
            True if the index should be (re)trained, False otherwise
        """
        # Check index type constraints
        if index_type.lower() in ["ivf", "ivfpq"] and is_trained:
            logger.debug(
                f"Index '{index_name}' type '{index_type}' cannot be retrained (already trained)"
            )
            return False

        # For ivfflat and ivfsq, allow retraining even if already trained
        # Use n_lists = 1 for untrained indexes as per requirement
        effective_n_lists = n_lists if is_trained else 1

        threshold = effective_n_lists * self._retrain_threshold
        should_train = num_vectors > threshold

        if should_train:
            action = "retraining" if is_trained else "training"
            logger.info(
                f"Index '{index_name}' (type: {index_type}) meets {action} criteria: "
                f"{num_vectors} vectors > {threshold} threshold "
                f"(n_lists={effective_n_lists}, trained={is_trained})"
            )

        return should_train

    def is_training(self, index_name: str) -> bool:
        """Check if specific index is currently training or queued."""
        with self._lock:
            if self._currently_training == index_name:
                return True
        # Check queue
        return self._is_in_queue(index_name)

    def _is_in_queue(self, index_name: str) -> bool:
        """Check if an index is in the training queue."""
        # Note: This is a somewhat expensive operation but queue should be small
        with self._lock:
            return any(job.index_name == index_name for job in list(self._queue.queue))

    def is_any_training(self) -> bool:
        """Check if any training is in progress."""
        with self._lock:
            return self._currently_training is not None

    def queue_training(self, job: TrainingJob) -> str:
        """
        Queue an index for training. Returns status.
        Used by auto-trigger after upsert.

        Args:
            job: TrainingJob with index details and parameters

        Returns:
            Status string: "queued", "already_queued", or "already_training"
        """
        if self.is_training(job.index_name):
            logger.info(
                f"Index '{job.index_name}' is already training or queued, skipping"
            )
            return "already_queued"

        self._queue.put(job)
        logger.info(f"Index '{job.index_name}' queued for training")
        self._ensure_worker_running()
        return "queued"

    def train_sync(self, job: TrainingJob, wait: bool = False) -> Dict[str, Any]:
        """
        Train an index. Used by /train endpoint.

        Args:
            job: TrainingJob with index details and parameters
            wait: If True, queues the job and blocks until training completes.
                  If False, queues and returns immediately.

        Returns:
            Dict with status and message
        """
        if not wait:
            # Check if already training/queued before adding duplicate
            with self._lock:
                if self._currently_training == job.index_name:
                    return {
                        "status": "in_progress",
                        "message": f"Index '{job.index_name}' is currently being trained",
                    }

            if self._is_in_queue(job.index_name):
                return {
                    "status": "queued",
                    "message": f"Index '{job.index_name}' is already queued for training",
                }

            self.queue_training(job)
            return {
                "status": "queued",
                "message": f"Index '{job.index_name}' queued for training",
            }

        # wait=True: Queue the job with a completion event and block until done.
        # This ensures that even if auto-training is already in progress,
        # this job will run after it finishes.
        job._completion_event = threading.Event()
        self._queue.put(job)
        logger.info(
            f"Index '{job.index_name}' queued for synchronous training (wait=True)"
        )
        self._ensure_worker_running()

        # Block until the worker signals completion
        job._completion_event.wait()

        if job._completion_error:
            return {
                "status": "error",
                "message": f"Training failed for '{job.index_name}': internal error during training",
            }

        return {
            "status": "success",
            "message": f"Index '{job.index_name}' trained successfully",
        }

    def trigger_training_if_needed(
        self,
        index_name: str,
        index_key_hex: str,
        num_vectors: int,
        n_lists: int,
        is_trained: bool,
        index_type: str,
    ) -> bool:
        """
        Check if training is needed and queue it if so.

        Args:
            index_name: Name of the index
            index_key_hex: Hex-encoded encryption key for the index
            num_vectors: Current number of vectors
            n_lists: Number of lists in index config
            is_trained: Whether index is currently trained
            index_type: Type of the index (ivf, ivfpq, ivfflat, ivfsq)

        Returns:
            True if training was queued, False otherwise
        """
        # Check if training is needed
        if not self.should_retrain(
            index_name, num_vectors, n_lists, is_trained, index_type
        ):
            return False

        # Check if already training or queued
        if self.is_training(index_name):
            logger.info(f"Index '{index_name}' is already training or queued, skipping")
            return False

        # Queue the training job, preserving current n_lists for retraining
        # Only preserve n_lists if > 1 (already trained); if n_lists <= 1, pass 0
        # to let cyborgdb-core compute optimal n_lists based on num_vectors
        job = TrainingJob(
            index_name=index_name,
            index_key_hex=index_key_hex,
            n_lists=n_lists if n_lists > 1 else 0,
        )
        status = self.queue_training(job)
        return status == "queued"

    def _ensure_worker_running(self):
        """Start background worker if not running."""
        with self._lock:
            if not self._worker_running:
                self._worker_running = True
                thread = threading.Thread(
                    target=self._process_queue, daemon=True, name="training-worker"
                )
                thread.start()
                logger.debug("Training worker thread started")

    def _process_queue(self):
        """Background worker that processes training queue."""
        logger.info("Training queue worker started")

        while True:
            try:
                # Wait for a job with timeout
                job = self._queue.get(timeout=5.0)
            except queue.Empty:
                # No jobs for 5 seconds, check if we should exit
                with self._lock:
                    if self._queue.empty():
                        self._worker_running = False
                        logger.info("Training queue worker exiting (queue empty)")
                        return
                continue

            # Process the job
            with self._lock:
                self._currently_training = job.index_name

            logger.info(f"Training worker processing index '{job.index_name}'")

            try:
                self._do_training(job)
                logger.info(
                    f"Training completed successfully for index '{job.index_name}'"
                )
            except Exception as e:
                logger.error(f"Training failed for '{job.index_name}': {str(e)}")
                job._completion_error = str(e)
            finally:
                with self._lock:
                    self._currently_training = None
                if job._completion_event:
                    job._completion_event.set()
                self._queue.task_done()

    def _do_training(self, job: TrainingJob):
        """
        Perform the actual training operation.

        Args:
            job: TrainingJob with index details and parameters
        """
        logger.info(f"Starting training for index '{job.index_name}'")

        if self._training_client is None:
            raise RuntimeError("Training client not initialized")

        # Load the index with the training client
        index_key = hex_to_bytes(job.index_key_hex)
        index = self._training_client.load_index(
            index_name=job.index_name, index_key=index_key
        )

        # Perform training
        # n_lists=0 means "compute optimal" in cyborgdb-core
        index.train(
            n_lists=job.n_lists if job.n_lists is not None else 0,
        )

        # Clear the index from the main service cache so subsequent requests
        # get the freshly trained index instead of a stale untrained version
        clear_index_cache(job.index_name)
        logger.info(f"Cleared index cache for '{job.index_name}' after training")

    def get_training_status(self) -> Dict[str, Any]:
        """Get the current training status."""
        with self._lock:
            queued_indexes = [job.index_name for job in list(self._queue.queue)]

            # Build training_indexes: currently training (if any) + queued indexes
            training_indexes = []
            if self._currently_training:
                training_indexes.append(self._currently_training)
            training_indexes.extend(queued_indexes)

            return {
                "training_indexes": training_indexes,
                "currently_training": self._currently_training,
                "queued_indexes": queued_indexes,
                "retrain_threshold": self._retrain_threshold,
                "worker_running": self._worker_running,
                # Backward compatibility fields for SDK
                "worker_pid": 0,
                "global_training": {},
            }

    def get_queue_position(self, index_name: str) -> Optional[int]:
        """
        Get the position of an index in the training queue.

        Args:
            index_name: Name of the index

        Returns:
            Position (0-indexed) or None if not in queue
        """
        with self._lock:
            if self._currently_training == index_name:
                return 0  # Currently being trained

            queue_list = list(self._queue.queue)
            for i, job in enumerate(queue_list):
                if job.index_name == index_name:
                    return i + 1  # +1 because currently_training is position 0
        return None

    def shutdown(self):
        """Shutdown the training manager and cleanup resources."""
        logger.info("Shutting down TrainingManager")

        # Wait for queue to be processed (with timeout)
        try:
            self._queue.join()
        except Exception:
            logger.exception(
                "Error while waiting for training queue to drain during shutdown; continuing shutdown"
            )


# Global training manager instance
_training_manager: Optional[TrainingManager] = None


def get_training_manager() -> TrainingManager:
    """Get or create the global training manager instance."""
    global _training_manager
    if _training_manager is None:
        _training_manager = TrainingManager()
    return _training_manager


def shutdown_training_manager():
    """Shutdown the global training manager."""
    global _training_manager
    if _training_manager is not None:
        _training_manager.shutdown()
        _training_manager = None
