# src/utils/llm_utils.py
from fastapi import HTTPException
from typing import Tuple

import asyncio
import logging

from src.llms.openai_llm import OpenAILanguageModel
from src.llms.ollama_llm import OllamaLanguageModel
from src.llms.bert_llm import BERTLanguageModel
from src.llms.falcon_llm import FalconLanguageModel
from src.llms.llama_llm import LlamaLanguageModel
from src.embeddings.huggingface_embedding import HuggingFaceEmbedding
from src.vectorstores.chroma_vectorstore import ChromaVectorStore
from src.vectorstores.chroma_manager import ChromaManager
from src.utils.logger import logger
from config.config import settings

# Global vector store instance for reuse
_vector_store = None
_embedding_model = None
_vs_lock = asyncio.Lock()


def get_llm_instance(provider: str):
    """
    Get LLM instance based on provider

    Args:
        provider (str): Name of the LLM provider

    Returns:
        BaseLLM: Instance of the LLM

    Raises:
        ValueError: If provider is not supported
    """
    llm_map = {
        'openai': lambda: OpenAILanguageModel(api_key=settings.OPENAI_API_KEY),
        'ollama': lambda: OllamaLanguageModel(base_url=settings.OLLAMA_BASE_URL),
        'bert': lambda: BERTLanguageModel(),
        'falcon': lambda: FalconLanguageModel(),
        'llama': lambda: LlamaLanguageModel(),
    }

    if provider not in llm_map:
        raise ValueError(f"Unsupported LLM provider: {provider}")
    return llm_map[provider]()


async def get_vector_store() -> Tuple[ChromaVectorStore, HuggingFaceEmbedding]:
    """
    Get vector store and embedding model instances with proper initialization

    Returns:
        Tuple[ChromaVectorStore, HuggingFaceEmbedding]: 
            Vector store and embedding model instances
    """
    global _vector_store, _embedding_model, _vs_lock

    async with _vs_lock:
        if _vector_store is not None and _embedding_model is not None:
            return _vector_store, _embedding_model

        try:
            # Load embedding model
            _embedding_model = HuggingFaceEmbedding(
                model_name=settings.EMBEDDING_MODEL)
            logger.info(f"Loaded embedding model: {settings.EMBEDDING_MODEL}")

            # Get ChromaDB client through the manager
            try:
                client = await ChromaManager.get_client(
                    persist_directory=settings.CHROMA_PATH,
                    reset_if_needed=True
                )
                logger.info("Successfully initialized ChromaDB client")
            except Exception as e:
                logger.error(f"Error getting ChromaDB client: {str(e)}")

                # Try to reset ChromaDB completely
                await ChromaManager.reset_chroma(settings.CHROMA_PATH)
                client = await ChromaManager.get_client(
                    persist_directory=settings.CHROMA_PATH
                )
                logger.info("Recreated ChromaDB client after reset")

            # Create and initialize vector store
            _vector_store = ChromaVectorStore(
                embedding_function=_embedding_model.embed_documents,
                persist_directory=settings.CHROMA_PATH,
                collection_name="documents",
                client=client
            )

            # Initialize the vector store
            await _vector_store.initialize()
            logger.info("Vector store successfully initialized")

            return _vector_store, _embedding_model

        except Exception as e:
            logger.error(f"Error initializing vector store: {str(e)}")
            raise HTTPException(
                status_code=500,
                detail=f"Failed to initialize vector store: {str(e)}"
            )


async def cleanup_vectorstore():
    """
    Cleanup and reset vector store resources
    """
    global _vector_store, _embedding_model, _vs_lock

    async with _vs_lock:
        _vector_store = None
        _embedding_model = None

        # Force garbage collection
        import gc
        gc.collect()

        # Reset ChromaDB completely
        await ChromaManager.reset_chroma(settings.CHROMA_PATH)