import hashlib
import threading
from typing import Dict, Tuple
from fastapi import HTTPException, status

# Attempt to import cyborgdb_core first
import cyborgdb_core as cyborgdb
from cyborgdb_core import Client, DBConfig, GPUConfig

from cyborgdb_service.core.config import settings
from cyborgdb_service.core.security import hex_to_bytes

# Thread-local storage for client instances
_thread_local = threading.local()

# Global index cache - stores (index, key_hash) tuples by index_name
# This cache is shared across all threads
_index_cache: Dict[str, Tuple[object, str]] = {}
_cache_lock = threading.Lock()


def _hash_key(key_hex: str) -> str:
    """Create a SHA-256 hash of the key for secure comparison."""
    return hashlib.sha256(key_hex.encode()).hexdigest()


def initialize_client():
    """
    Initialize the CyborgDB client.
    """

    try:
        # Initialize client with settings from config
        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)

        client = Client(
            api_key=settings.CYBORGDB_API_KEY,
            index_location=index_location,
            config_location=config_location,
            items_location=items_location,
            cpu_threads=settings.CPU_THREADS,
            gpu_config=gpu_config,
        )

        # Print GPU configuration
        print(
            f"Main client initialized - GPU enabled: {client.is_gpu_enabled()}, "
            f"GPU config: {client.get_gpu_config()}"
        )

        return client
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to initialize CyborgDB client: {str(e)}",
        )


def get_client():
    """
    Get the CyborgDB client instance. Thread-safe implementation.
    """
    # Check if the current thread has a client
    if not hasattr(_thread_local, "client"):
        # If not, create one for this thread
        _thread_local.client = initialize_client()

    return _thread_local.client


def load_index(index_name: str, index_key_hex: str):
    """
    Load an index with the provided encryption key. Uses caching to avoid reloading.
    Validates the encryption key hash before returning cached indexes.
    """
    key_hash = _hash_key(index_key_hex)
    # Check cache first (using only index_name as key)
    with _cache_lock:
        cached_entry = _index_cache.get(index_name)
        if cached_entry is not None:
            cached_index, cached_key_hash = cached_entry
            if cached_key_hash == key_hash:
                return cached_index
            # Key mismatch - reject with auth error
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Invalid encryption key for index",
            )

    # Not in cache, need to load from core
    current_client = get_client()
    index_key = hex_to_bytes(index_key_hex)

    try:
        index = current_client.load_index(index_name=index_name, index_key=index_key)

        # Add to cache with key hash
        index_key = 0
        index_key_hex = 0
        with _cache_lock:
            _index_cache[index_name] = (index, key_hash)

        return index
    except ValueError as e:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail=f"Index not found: {str(e)}"
        )
    except RuntimeError as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to load index: {str(e)}",
        )


def clear_index_cache(index_name: str = None):
    """
    Clear the index cache. Can clear a specific index or all indexes.

    Args:
        index_name: If provided, only clear that specific index from cache
    """
    with _cache_lock:
        if index_name:
            if index_name in _index_cache:
                del _index_cache[index_name]
        else:
            _index_cache.clear()


def create_index_config(config_data):
    """
    Create the appropriate IndexConfig object from request data.
    """
    try:
        if config_data.type == "ivf":
            return cyborgdb.IndexIVF(dimension=config_data.dimension)
        elif config_data.type == "ivfpq":
            return cyborgdb.IndexIVFPQ(
                dimension=config_data.dimension,
                pq_dim=config_data.pq_dim,
                pq_bits=config_data.pq_bits,
            )
        elif config_data.type == "ivfflat":
            return cyborgdb.IndexIVFFlat(dimension=config_data.dimension)
        elif config_data.type == "ivfsq":
            return cyborgdb.IndexIVFSQ(
                dimension=config_data.dimension,
                sq_bits=config_data.sq_bits,
            )
        else:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Unsupported index type: {config_data.type}",
            )
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Failed to create index config: {str(e)}",
        )
