from pinecone import Pinecone, ServerlessSpec
import os
from dotenv import load_dotenv

load_dotenv() 

class PineconeStore:
    def __init__(self):
        self.pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))

        if 'blip-soccer-highlights' not in self.pc.list_indexes().names():
            self.pc.create_index(
                name='blip-soccer-highlights',
                dimension=1024,  # This must match InternVL2_5-8B-MPO
                metric='cosine',
                spec=ServerlessSpec(cloud='aws', region=os.getenv("PINECONE_ENV", "us-east-1"))
            )

        self.index = self.pc.Index('blip-soccer-highlights')

    def upsert(self, id, vector, metadata):
        self.index.upsert([(id, vector.tolist(), metadata)])

    def query(self, vector, filter_key, top_k):
        results = self.index.query(vector.tolist(), top_k=top_k, include_metadata=True)
        return [
            match["metadata"] for match in results["matches"]
            if filter_key in match["id"]
        ]