import os
import hashlib
from typing import List
import pypdf
from docx import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from .config import settings

class DocumentProcessor:
    def __init__(self):
        self.embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
        self.vector_store = Chroma(
            persist_directory=settings.CHROMA_DB_PATH,
            embedding_function=self.embeddings
        )
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=settings.CHUNK_SIZE,
            chunk_overlap=settings.CHUNK_OVERLAP
        )

    def get_file_checksum(self, file_path: str) -> str:
        with open(file_path, "rb") as f:
            return hashlib.md5(f.read()).hexdigest()

    def extract_text(self, file_path: str) -> str:
        ext = os.path.splitext(file_path)[1].lower()
        text = ""
        if ext == ".pdf":
            with open(file_path, "rb") as f:
                reader = pypdf.PdfReader(f)
                for page in reader.pages:
                    text += page.extract_text() + "\n"
        elif ext == ".docx":
            doc = Document(file_path)
            for para in doc.paragraphs:
                text += para.text + "\n"
        elif ext == ".txt":
            with open(file_path, "r", encoding="utf-8") as f:
                text = f.read()
        return text

    def ingest_document(self, file_path: str, metadata: dict = None):
        text = self.extract_text(file_path)
        if not text.strip():
            return None
        
        chunks = self.text_splitter.split_text(text)
        metadatas = [metadata or {} for _ in chunks]
        for meta in metadatas:
            meta["source"] = os.path.basename(file_path)
            meta["checksum"] = self.get_file_checksum(file_path)
            
        self.vector_store.add_texts(texts=chunks, metadatas=metadatas)
        return len(chunks)

    def search(self, query: str, k: int = 5):
        return self.vector_store.similarity_search(query, k=k)

processor = DocumentProcessor()