import pandas as pd
from utils.vector_database import search_in_milvus, fashionclip_collection, fashionsiglip_collection
from utils.embedding_generation import generate_query_embedding
from utils.load_models import fclip_model, fclip_processor
from utils.load_models import siglip_model, siglip_preprocess_val, siglip_tokenizer

# Function to dynamically select the Milvus collection and search field
def get_milvus_collection_and_field(model_type, embedding_type):
    # Define mapping of model and embedding types to collections and fields
    if model_type == "fashionCLIP":
        collection = fashionclip_collection
        if embedding_type == "text":
            search_field = "text_embedding"
        elif embedding_type == "image":
            search_field = "image_embedding"
        elif embedding_type == "average":
            search_field = "avg_embedding"
        elif embedding_type == "weighted average":
            search_field = "weighted_avg_embedding"
    elif model_type == "fashionSigLIP":
        collection = fashionsiglip_collection
        if embedding_type == "text":
            search_field = "text_embedding"
        elif embedding_type == "image":
            search_field = "image_embedding"
        elif embedding_type == "average":
            search_field = "avg_embedding"
        elif embedding_type == "weighted average":
            search_field = "weighted_avg_embedding"
    else:
        raise ValueError("Invalid model type. Choose 'fashionCLIP' or 'fashionSigLIP'.")

    return collection, search_field

# Function to handle the complete search flow
def search(query, query_type, model_type, embedding_type):
    # Step 1: Generate the query embedding based on the user input and model type
    if model_type == "fashionCLIP":
        query_embedding = generate_query_embedding(query, query_type, fclip_model, fclip_processor, fclip_processor, "fashionCLIP")
    elif model_type == "fashionSigLIP":
        query_embedding = generate_query_embedding(query, query_type, siglip_model, siglip_preprocess_val, siglip_tokenizer, "fashionSigLIP")

    # Step 2: Get the appropriate Milvus collection and search field
    collection, search_field = get_milvus_collection_and_field(model_type, embedding_type)

    # Step 3: Perform search in Milvus using the query embedding
    search_results = search_in_milvus(collection, search_field, query_embedding, top_k=10)

    # Step 4: Extract images, similarity scores, and metadata from the search results
    images = [result['image'] for result in search_results]
    scores = [result['similarity_score'] for result in search_results]
    metadata = [result['metadata'] for result in search_results]

    return images, scores, metadata

# Function to run the search and get results for both models
def run_search(query_type, embedding_type, query_input_text, query_input_image):
    if query_type == "text":
        query = query_input_text
    else:
        query = query_input_image

    # Perform search for FashionCLIP
    fclip_images, fclip_scores, fclip_metadata = search(query, query_type, "fashionCLIP", embedding_type)

    # Perform search for MARGO-FashionSigLip
    siglip_images, siglip_scores, siglip_metadata = search(query, query_type, "fashionSigLIP", embedding_type)

    # Convert scores and metadata into a pandas DataFrame for each model
    fclip_results_df = pd.DataFrame({
        "Score": fclip_scores,
        "Metadata": fclip_metadata,
    })

    siglip_results_df = pd.DataFrame({
        "Score": siglip_scores,
        "Metadata": siglip_metadata,
    })

    return fclip_images, fclip_results_df, siglip_images, siglip_results_df