"""
Pinecone vector store wrapper.
Reconstructed to match the interface app.py expects:
  - create_vector_store_from_env()
  - .add_documents(documents, chunk_size, chunk_overlap)
  - .get_index_stats()
  - .similarity_search(...) / .similarity_search_with_score(...)
  - .delete_index()
"""

import os
import time
from typing import List, Tuple

from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from pinecone import Pinecone, ServerlessSpec

# text-embedding-3-small / ada-002 both emit 1536 dims
EMBEDDING_DIMENSION = 1536


class VectorStore:
    """Thin wrapper around a Pinecone serverless index + LangChain embeddings."""

    def __init__(
        self,
        api_key: str,
        index_name: str,
        embedding_model: str = "text-embedding-3-small",
        cloud: str = "aws",
        region: str = "us-east-1",
    ):
        self.index_name = index_name
        self.pc = Pinecone(api_key=api_key)
        self.embeddings = OpenAIEmbeddings(model=embedding_model)

        existing = [idx["name"] for idx in self.pc.list_indexes()]
        if index_name not in existing:
            self.pc.create_index(
                name=index_name,
                dimension=EMBEDDING_DIMENSION,
                metric="cosine",
                spec=ServerlessSpec(cloud=cloud, region=region),
            )
            # wait for the index to be ready
            while not self.pc.describe_index(index_name).status["ready"]:
                time.sleep(1)

        self.index = self.pc.Index(index_name)
        self._store = PineconeVectorStore(
            index=self.index, embedding=self.embeddings
        )

    def add_documents(
        self,
        documents: List[Document],
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
    ) -> int:
        """Split documents into chunks and upsert them. Returns chunk count."""
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=["\n\n", "\n", ". ", " ", ""],
        )
        chunks = splitter.split_documents(documents)
        if chunks:
            self._store.add_documents(chunks)
        return len(chunks)

    def similarity_search(self, query: str, k: int = 5) -> List[Document]:
        return self._store.similarity_search(query, k=k)

    def similarity_search_with_score(
        self, query: str, k: int = 5
    ) -> List[Tuple[Document, float]]:
        return self._store.similarity_search_with_score(query, k=k)

    def as_retriever(self, k: int = 5):
        return self._store.as_retriever(search_kwargs={"k": k})

    def get_index_stats(self) -> dict:
        stats = self.index.describe_index_stats()
        return {
            "total_vector_count": stats.get("total_vector_count", 0),
            "dimension": stats.get("dimension", EMBEDDING_DIMENSION),
        }

    def delete_index(self) -> None:
        """Destroy the entire index. Irreversible."""
        self.pc.delete_index(self.index_name)


def create_vector_store_from_env() -> VectorStore:
    """Build a VectorStore from environment variables."""
    api_key = os.getenv("PINECONE_API_KEY")
    index_name = os.getenv("PINECONE_INDEX_NAME", "rag-demo")
    region = os.getenv("PINECONE_ENVIRONMENT", "us-east-1")
    embedding_model = os.getenv(
        "OPENAI_EMBEDDING_MODEL", "text-embedding-3-small"
    )

    if not api_key:
        raise ValueError("PINECONE_API_KEY is not set")

    # PINECONE_ENVIRONMENT may arrive as "us-east-1-aws"; keep just the region
    region = region.replace("-aws", "").replace("-gcp", "")

    return VectorStore(
        api_key=api_key,
        index_name=index_name,
        embedding_model=embedding_model,
        region=region,
    )