from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
import os
from typing import List, Dict
from .embeddings import EmbeddingService

class PDFProcessor:
    def __init__(self, embedding_service: EmbeddingService):
        self.embedding_service = embedding_service
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50,
            length_function=len,
        )

    def load_and_split(self, file_path: str, metadata: Dict) -> List[Dict]:
        """Load and split a PDF file."""
        loader = PyPDFLoader(file_path)
        documents = loader.load()

        chunks = self.text_splitter.split_documents(documents)

        # Convert to dictionaries with metadata
        chunk_data = []
        for i, doc in enumerate(chunks):
            chunk_data.append({
                "content": doc.page_content,
                "metadata": {
                    **metadata,
                    "chunk_id": i,
                    "source": file_path,
                }
            })

        return chunk_data

    def embed_chunks(self, chunks: List[Dict]) -> List[Dict]:
        """Add embeddings to chunks."""
        texts = [chunk["content"] for chunk in chunks]
        embeddings = self.embedding_service.embed(texts)

        for chunk, embedding in zip(chunks, embeddings):
            chunk["embedding"] = embedding.tolist()

        return chunks