import pinecone
from typing import List, Dict, Any, Optional
import uuid
from config import settings

class VectorStore:
    def __init__(self):
        # Initialize Pinecone
        pinecone.init(
            api_key=settings.PINECONE_API_KEY,
            environment=settings.PINECONE_ENVIRONMENT
        )
        
        # Create or get the index
        if settings.PINECONE_INDEX_NAME not in pinecone.list_indexes():
            pinecone.create_index(
                name=settings.PINECONE_INDEX_NAME,
                dimension=settings.VECTOR_DIMENSION,
                metric="cosine"
            )
        
        self.index = pinecone.Index(settings.PINECONE_INDEX_NAME)
    
    def add_content(self, content: str, metadata: Optional[Dict[str, Any]] = None) -> str:
        """
        Add content to the vector store with optional metadata.
        Returns the ID of the added content.
        """
        content_id = str(uuid.uuid4())
        
        # Prepare metadata
        if metadata is None:
            metadata = {}
        metadata['content'] = content
        
        # Upsert the vector with metadata
        self.index.upsert(
            vectors=[(content_id, [0] * settings.VECTOR_DIMENSION, metadata)],
            namespace="content"
        )
        
        return content_id
    
    def search(self, query_vector: List[float], top_k: int = settings.MAX_SEARCH_RESULTS) -> List[Dict[str, Any]]:
        """
        Search for similar content using a query vector.
        Returns a list of dictionaries containing content and metadata.
        """
        results = self.index.query(
            vector=query_vector,
            top_k=top_k,
            include_metadata=True,
            namespace="content"
        )
        
        return [
            {
                'id': match.id,
                'content': match.metadata['content'],
                'score': match.score,
                **{k: v for k, v in match.metadata.items() if k != 'content'}
            }
            for match in results.matches
        ]
    
    def get_all_content(self) -> List[Dict[str, Any]]:
        """
        Retrieve all content from the vector store.
        """
        # Fetch all vectors from the index
        results = self.index.query(
            vector=[0] * settings.VECTOR_DIMENSION,
            top_k=10000,  # Adjust based on your needs
            include_metadata=True,
            namespace="content"
        )
        
        return [
            {
                'id': match.id,
                'content': match.metadata['content'],
                **{k: v for k, v in match.metadata.items() if k != 'content'}
            }
            for match in results.matches
        ]
    
    def delete_content(self, content_id: str) -> bool:
        """
        Delete content from the vector store by ID.
        Returns True if successful, False otherwise.
        """
        try:
            self.index.delete(ids=[content_id], namespace="content")
            return True
        except Exception as e:
            print(f"Error deleting content: {e}")
            return False 