import os
import fitz
import spacy
import matplotlib.pyplot as plt
from langchain.docstore.document import Document
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings

# Load SpaCy model once
nlp = spacy.load("en_core_web_sm")

def spacy_sentence_tokenize(text: str) -> list:
    doc = nlp(text)
    return [sent.text.strip() for sent in doc.sents]

def load_pdf_to_documents(pdf_path: str) -> list:
    documents = []
    with fitz.open(pdf_path) as doc:
        for i, page in enumerate(doc):
            text = page.get_text().replace("-\n", "").replace("\n", " ").strip()
            if text:
                documents.append(Document(page_content=text, metadata={"page": i}))
    return documents

def sentence_overlap_chunk(text: str, max_tokens: int = 150, overlap_sent_count: int = 2) -> list:
    sentences = spacy_sentence_tokenize(text)
    chunks, current_chunk, current_len = [], [], 0

    for sentence in sentences:
        token_count = len(sentence.split())
        if current_len + token_count <= max_tokens:
            current_chunk.append(sentence)
            current_len += token_count
        else:
            chunks.append(" ".join(current_chunk))
            current_chunk = current_chunk[-overlap_sent_count:] + [sentence]
            current_len = sum(len(s.split()) for s in current_chunk)

    if current_chunk:
        chunks.append(" ".join(current_chunk))
    return chunks

def analyze_chunks(chunks: list):
    token_lengths = [len(chunk.page_content.split()) for chunk in chunks]

    print(f"Total Chunks: {len(token_lengths)}")
    print(f"Avg Tokens per Chunk: {sum(token_lengths)/len(token_lengths):.2f}")
    print(f"Min Tokens: {min(token_lengths)}")
    print(f"Max Tokens: {max(token_lengths)}")

    plt.hist(token_lengths, bins=20)
    plt.title("Chunk Token Length Distribution")
    plt.xlabel("Token Count")
    plt.ylabel("Number of Chunks")
    plt.show()

def build_vector_store(documents: list,
                       max_tokens: int = 250,
                       overlap_sent_count: int = 3,
                       model_name: str = "sentence-transformers/allenai-specter",
                       persist_directory: str = "./vector_db") -> FAISS:

    all_chunks = []
    for doc in documents:
        chunks = sentence_overlap_chunk(doc.page_content, max_tokens=max_tokens, overlap_sent_count=overlap_sent_count)
        all_chunks.extend([Document(page_content=chunk, metadata=doc.metadata) for chunk in chunks])

    analyze_chunks(all_chunks)

    embeddings = HuggingFaceEmbeddings(model_name=model_name)
    vectorstore = FAISS.from_documents(all_chunks, embeddings)
    vectorstore.save_local(persist_directory)
    return vectorstore

def load_vector_store(persist_directory: str,
                      model_name: str = "sentence-transformers/allenai-specter") -> FAISS:
    embeddings = HuggingFaceEmbeddings(model_name=model_name)
    return FAISS.load_local(persist_directory, embeddings, allow_dangerous_deserialization=True)

def query_vector_store(vectorstore: FAISS, query: str, k: int = 3, show: bool = True):
    results = vectorstore.similarity_search(query, k=k)
    if show:
        for i, doc in enumerate(results, 1):
            print(f"\n--- Result {i} (Page {doc.metadata.get('page')}):\n{doc.page_content[:500]}...\n")
    return results