# embedding_manager.py
import os
from typing import Optional, List
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings  # modern import
from langchain_community.vectorstores import Chroma

class EmbeddingManager:
    """
    Extracts text from a PDF and builds/reuses a persisted Chroma vector store.
    Persistence path: ./embeddings/<PDF_STEM>
    """
    def __init__(
        self,
        pdf_path: str,
        base_dir: str = "./embeddings",
        chunk_size: int = 512,
        chunk_overlap: int = 64,
        embedding_model: str = "text-embedding-3-large",
        openai_api_key_env: str = "OPENAI_API_KEY",
    ):
        self.pdf_path = pdf_path
        self.base_dir = base_dir
        name = os.path.splitext(os.path.basename(pdf_path))[0]
        self.txt_path = os.path.join(base_dir, f"{name}.txt")
        self.persist_dir = os.path.join(base_dir, name)
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.embedding_model = embedding_model
        self.openai_api_key_env = openai_api_key_env
        os.makedirs(base_dir, exist_ok=True)

    def pdf_to_txt(self) -> str:
        """Dump PDF text to a .txt (idempotent)."""
        if os.path.exists(self.txt_path):
            print(f"[INFO] Using existing text at {self.txt_path}")
            return self.txt_path

        reader = PdfReader(self.pdf_path)
        with open(self.txt_path, "w", encoding="utf-8") as f:
            for page in reader.pages:
                text = page.extract_text()
                if text:
                    f.write(text + "\n")
        print(f"[INFO] Extracted text to {self.txt_path}")
        return self.txt_path

    def _load_embeddings(self) -> OpenAIEmbeddings:
        key = os.environ.get(self.openai_api_key_env)
        if not key:
            raise RuntimeError(
                f"Missing {self.openai_api_key_env} in environment. "
                "Set it in your Hugging Face Space → Settings → Secrets."
            )
        # Modern LangChain uses langchain_openai.OpenAIEmbeddings
        return OpenAIEmbeddings(api_key=key, model=self.embedding_model)

    def get_or_create_embeddings(self) -> Chroma:
        """
        Returns a Chroma vector store, creating & persisting if needed.
        """
        embeddings = self._load_embeddings()

        if os.path.exists(self.persist_dir) and os.listdir(self.persist_dir):
            print(f"[INFO] Loading embeddings from {self.persist_dir}")
            return Chroma(persist_directory=self.persist_dir, embedding_function=embeddings)

        txt = self.pdf_to_txt()
        with open(txt, "r", encoding="utf-8") as f:
            text = f.read()
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
        )
        chunks: List[str] = splitter.split_text(text)

        vectordb = Chroma.from_texts(
            chunks,
            embedding=embeddings,
            persist_directory=self.persist_dir
        )
        vectordb.persist()
        print(f"[INFO] Created embeddings in {self.persist_dir}")
        return vectordb