import torch
from pinecone import Pinecone
from torchvision.transforms import Compose, Resize, ToTensor, Normalize
from PIL import Image
from transformers import CLIPProcessor, CLIPModel

# Initialize the Pinecone client
pc = Pinecone(
    api_key="pcsk_6w79Cx_MfeLRKxmSKerQhKiyrjZieguMKdiSVV1M3no8fTzAygLT8Qpx8F7AS7iMSzUEX1",
    environment="us-east-1"  # Ensure this matches your region
)

# Connect to the existing index
index_name = "multimodal-search"
if index_name not in pc.list_indexes().names():
    raise ValueError(f"Index {index_name} does not exist. Please create it in your Pinecone dashboard.")

index = pc.Index(index_name)

# Initialize CLIP model and processor
device = "cuda" if torch.cuda.is_available() else "cpu"
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Define global image transformation pipeline
transform = Compose([
    Resize((224, 224)),
    ToTensor(),
    Normalize(mean=(0.481, 0.457, 0.408), std=(0.268, 0.261, 0.275))  # CLIP-specific normalization
])

def query_pinecone_with_filters(query, mode="text-to-image", top_k=5, filter_metadata=None):
    """
    Perform text-to-image or image-to-image search with metadata filtering using Pinecone and CLIP.
    Args:
        query: String (for text-to-image) or file path to image (for image-to-image).
        mode: "text-to-image" or "image-to-image".
        top_k: Number of top results to retrieve.
        filter_metadata: Dictionary of filters to apply on metadata (e.g., {"type": "image", "captions": "dog"}).
    Returns:
        List of matches with metadata.
    """
    if mode == "text-to-image":
        # Convert text query to embedding
        text_inputs = processor(text=[query], return_tensors="pt", padding=True).to(device)
        query_embedding = model.get_text_features(**text_inputs)
        query_embedding = query_embedding / query_embedding.norm(dim=-1, keepdim=True)  # Normalize

    elif mode == "image-to-image":
        # Convert image to embedding
        query_image = Image.open(query).convert("RGB")
        image_tensor = transform(query_image).unsqueeze(0).to(device)
        query_embedding = model.get_image_features(image_tensor)
        query_embedding = query_embedding / query_embedding.norm(dim=-1, keepdim=True)  # Normalize

    else:
        raise ValueError("Invalid mode. Choose 'text-to-image' or 'image-to-image'.")

    # Query Pinecone
    query_vector = query_embedding.cpu().detach().numpy().flatten().tolist()
    results = index.query(
        vector=query_vector,
        top_k=top_k,
        include_metadata=True,
        filter=filter_metadata
    )

    # Update image paths to use URLs
    for match in results["matches"]:
        image_filename = match["metadata"]["image_path"].split("\\")[-1]  # Extract filename
        match["metadata"]["image_path"] = f"https://Dev24910-VisionVerse-Search.hf.space/file/workspace/images/{image_filename}"

    return results["matches"]