from .config import Config
from pinecone import Pinecone, ServerlessSpec
from typing import List, Optional
import time


class VectorStore:
    def __init__(self):
        if Config.VECTOR_STORE_TYPE == "pinecone":
            self.pc = Pinecone(api_key=Config.PINECONE_API_KEY)

            # Free tier supported regions
            FREE_TIER_SUPPORTED_REGIONS = {
                'aws': 'us-east-1',
                'gcp': 'us-central1'
            }

            # Check if index exists
            if Config.PINECONE_INDEX not in self.pc.list_indexes().names():
                print(f"Creating new Pinecone index: {Config.PINECONE_INDEX}")
                try:
                    # First try AWS free tier region
                    self.pc.create_index(
                        name=Config.PINECONE_INDEX,
                        dimension=1024,  # Cohere embed-english-v3.0 dimension
                        metric="cosine",
                        spec=ServerlessSpec(
                            cloud="aws",
                            region=FREE_TIER_SUPPORTED_REGIONS['aws']
                        )
                    )
                except Exception as e:
                    print(f"AWS region failed, trying GCP: {str(e)}")
                    # Fallback to GCP if AWS fails
                    self.pc.create_index(
                        name=Config.PINECONE_INDEX,
                        dimension=1024,
                        metric="cosine",
                        spec=ServerlessSpec(
                            cloud="gcp",
                            region=FREE_TIER_SUPPORTED_REGIONS['gcp']
                        )
                    )

                # Wait for index to initialize
                time.sleep(1)

            self.index = self.pc.Index(Config.PINECONE_INDEX)

    def upsert_document(self, doc_id: str, embedding: List[float], metadata: dict):
        self.index.upsert(
            vectors=[{
                "id": doc_id,
                "values": embedding,
                "metadata": metadata
            }]
        )

    def search_similar(self, embedding: List[float], top_k: int = 3):
        return self.index.query(
            vector=embedding,
            top_k=top_k,
            include_metadata=True
        )

    def get_document(self, doc_id: str) -> Optional[dict]:
        fetch_response = self.index.fetch(ids=[doc_id])
        if doc_id in fetch_response.vectors:
            return fetch_response.vectors[doc_id]
        return None