"""
Vector storage configuration for the RAG solution using LlamaIndex and Qdrant.

This module provides initialization and configuration for:
- Qdrant vector storage connection
- Ollama embedding model
- Automatic collection creation
"""

import os
from typing import Optional
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.llms.ollama import Ollama
from qdrant_client import QdrantClient
from loguru import logger


def initialize_vector_storage(
    collection_name: str = "documents_llamaindex",
    host: str = "localhost",
    port: int = 6333,
    grpc_port: int = 6334,
    ollama_base_url: str = "http://localhost:11434",
    ollama_embed_model: Optional[str] = None
) -> tuple[QdrantVectorStore, VectorStoreIndex]:
    """
    Initialize Qdrant vector storage with Ollama embeddings.

    Args:
        collection_name: Name of the Qdrant collection
        host: Qdrant host address
        port: Qdrant REST API port
        grpc_port: Qdrant gRPC API port
        ollama_base_url: Base URL for Ollama API
        ollama_embed_model: Name of the Ollama embedding model

    Returns:
        Tuple of (QdrantVectorStore, VectorStoreIndex)
    """
    logger.info(f"Initializing vector storage with collection: {collection_name}")

    # Get embedding model from environment if not provided
    if ollama_embed_model is None:
        ollama_embed_model = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")

    logger.info(f"Using Ollama embedding model: {ollama_embed_model}")

    try:
        # Initialize Qdrant client
        client = QdrantClient(host=host, port=port)

        # Check if collection exists, create if it doesn't
        collections = client.get_collections().collections
        collection_names = [coll.name for coll in collections]

        if collection_name not in collection_names:
            logger.info(f"Collection '{collection_name}' does not exist, creating...")
            client.create_collection(
                collection_name=collection_name,
                vectors_config={
                    "size": 4096,  # Default size for most embedding models
                    "distance": "Cosine"  # Cosine distance is commonly used
                }
            )
            logger.info(f"Collection '{collection_name}' created successfully")
        else:
            logger.info(f"Collection '{collection_name}' already exists")

        # Initialize the Qdrant vector store
        vector_store = QdrantVectorStore(
            client=client,
            collection_name=collection_name
        )

        # Initialize Ollama embedding
        embed_model = OllamaEmbedding(
            model_name=ollama_embed_model,
            base_url=ollama_base_url
        )

        # Create index from vector store with the embedding model
        index = VectorStoreIndex.from_vector_store(
            vector_store=vector_store,
            embed_model=embed_model
        )

        logger.info("Vector storage initialized successfully")
        return vector_store, index

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


# Optional: Alternative embedding configuration using OpenAI via OpenRouter
# Uncomment and configure as needed for future use
# from llama_index.embeddings.openai import OpenAIEmbedding
#
# def initialize_openai_embeddings():
#     # Use OpenRouter API key from environment
#     os.environ["OPENAI_API_KEY"] = os.getenv("OPENROUTER_API_KEY", "")
#
#     embed_model = OpenAIEmbedding(
#         model="openai/text-embedding-3-small",  # Or another suitable model
#         api_base="https://openrouter.ai/api/v1"  # OpenRouter endpoint
#     )
#     return embed_model


def get_vector_store_and_index() -> tuple[QdrantVectorStore, VectorStoreIndex]:
    """
    Convenience function to get the initialized vector store and index.

    Returns:
        Tuple of (QdrantVectorStore, VectorStoreIndex)
    """
    return initialize_vector_storage()


if __name__ == "__main__":
    # Example usage
    logger.info("Testing vector storage initialization...")
    try:
        vector_store, index = get_vector_store_and_index()
        logger.info("Vector storage test successful!")
    except Exception as e:
        logger.error(f"Vector storage test failed: {e}")