from fastapi import APIRouter, Depends, Body
import logging

from cyborgdb_service.api.deps import get_current_client
from cyborgdb_service.core.security import hex_to_bytes
from cyborgdb_service.db.client import (
    create_index_config,
    load_index,
    clear_index_cache,
    _index_cache,
    _cache_lock,
    _hash_key,
)
from cyborgdb_service.core.training_manager import get_training_manager, TrainingJob
from cyborgdb_service.api.schemas.index import (
    CreateIndexRequest,
    TrainRequest,
    IndexOperationRequest,
    IndexInfoResponse,
    IndexListResponse,
    IndexTrainingStatusResponse,
    CreateResponses,
    TrainResponses,
    DeleteIndexResponses,
)
from cyborgdb_service.utils.error_handler import handle_exception

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/indexes")


@router.get("/list", summary="List Encrypted Indexes", responses=IndexListResponse)
async def list_indexes(client=Depends(get_current_client)):
    """
    List all available indexes.
    """
    try:
        indexes = client.list_indexes()
        return {"indexes": indexes}
    except Exception as e:
        handle_exception(e, "Failed to list indexes")


@router.post("/create", summary="Create Encrypted Index", responses=CreateResponses)
async def create_index(request: CreateIndexRequest, client=Depends(get_current_client)):
    """
    Create a new encrypted index with the provided configuration.
    """
    # Convert hex key to bytes
    index_key = hex_to_bytes(request.index_key)

    # Create the appropriate index config
    index_config = create_index_config(request.index_config)

    try:
        # Check for embedding dependencies if embedding_model is provided
        if request.embedding_model:
            try:
                import sentence_transformers  # noqa: F401
            except ImportError as e:
                raise ValueError(
                    "Embedding model requires sentence-transformers. "
                    "Install with: pip install 'cyborgdb-service[embeddings]'"
                ) from e

        # Create the index and cache it
        index = client.create_index(
            index_name=request.index_name,
            index_key=index_key,
            index_config=index_config,
            embedding_model=request.embedding_model,
            metric=request.metric,
        )

        index_key = 0
        # Cache the newly created index with key hash to avoid reloading
        with _cache_lock:
            _index_cache[request.index_name] = (index, _hash_key(request.index_key))

        return {
            "status": "success",
            "message": f"Index '{request.index_name}' created successfully",
        }
    except Exception as e:
        handle_exception(e, "Failed to create index")


@router.post(
    "/describe", summary="Describe Encrypted Index", responses=IndexInfoResponse
)
async def get_index_info(
    request: IndexOperationRequest = Body(...), client=Depends(get_current_client)
):
    """
    Get information about a specific index.
    """
    index = load_index(request.index_name, request.index_key)

    return {
        "index_name": index.index_name(),
        "index_type": index.index_type(),
        "is_trained": index.is_trained(),
        "index_config": index.index_config(),
    }


@router.post(
    "/delete", summary="Delete Encrypted Index", responses=DeleteIndexResponses
)
async def delete_index(
    request: IndexOperationRequest = Body(...), client=Depends(get_current_client)
):
    """
    Delete a specific index.
    """
    index = load_index(request.index_name, request.index_key)

    try:
        index.delete_index()

        # Clear the index from cache after successful deletion
        clear_index_cache(request.index_name)

        return {
            "status": "success",
            "message": f"Index '{request.index_name}' deleted successfully",
        }
    except Exception as e:
        handle_exception(e, "Failed to delete index")


@router.post("/train", summary="Train Encrypted index", responses=TrainResponses)
async def train_index(request: TrainRequest, client=Depends(get_current_client)):
    """
    Train the index for efficient querying.

    Training is queued and processed asynchronously. If the index is already
    being trained or queued, returns the current status.

    Note: Automatic training is also triggered based on vector count thresholds
    after upserts.
    """
    # Validate index exists before queueing
    load_index(request.index_name, request.index_key)

    training_manager = get_training_manager()

    # Build job with request values, falling back to TrainingJob defaults for None values
    job_kwargs: dict = {
        "index_name": request.index_name,
        "index_key_hex": request.index_key,
    }
    if request.n_lists is not None:
        job_kwargs["n_lists"] = request.n_lists
    if request.batch_size is not None:
        job_kwargs["batch_size"] = request.batch_size
    if request.max_iters is not None:
        job_kwargs["max_iters"] = request.max_iters
    if request.tolerance is not None:
        job_kwargs["tolerance"] = request.tolerance
    if request.max_memory is not None:
        job_kwargs["max_memory"] = request.max_memory

    job = TrainingJob(**job_kwargs)
    result = training_manager.train_sync(job, wait=True)
    return result


@router.get(
    "/training-status",
    summary="Get Training Status",
    responses=IndexTrainingStatusResponse,
)
async def get_training_status():
    """
    Get the current training status including indexes being trained
    and the retrain threshold configuration.

    Returns:
        dict: Training status information including:
            - training_indexes: List of index names currently being trained
            - retrain_threshold: The multiplier used for the retraining threshold
    """
    try:
        training_manager = get_training_manager()
        return training_manager.get_training_status()
    except Exception as e:
        handle_exception(e, "Failed to get training status")
