# -*- coding: utf-8 -*-
""" 
@title: PDF AI Assistant
@author: Your Name
"""
# app.py
import streamlit as st
import fitz  # PyMuPDF
import numpy as np
import torch
from sentence_transformers import SentenceTransformer, util
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from tqdm.auto import tqdm
import textwrap
import re
import pandas as pd

# Configuration
MODEL_NAME = "all-mpnet-base-v2"
LLM_MODEL_ID = "google/gemma-2b-it"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
NUM_RESULTS = 5
MIN_TOKEN_LENGTH = 30

# Initialize session state
if 'processed' not in st.session_state:
    st.session_state.processed = False
if 'embeddings' not in st.session_state:
    st.session_state.embeddings = None
if 'pages_and_chunks' not in st.session_state:
    st.session_state.pages_and_chunks = []

# Helper functions
def text_formatter(text: str) -> str:
    cleaned_text = text.replace("\n", " ").strip()
    cleaned_text = re.sub(r'\s+', ' ', cleaned_text)
    return cleaned_text

def split_list(input_list: list, slice_size: int = 10) -> list[list[str]]:
    return [input_list[i:i + slice_size] for i in range(0, len(input_list), slice_size)]

def print_wrapped(text, wrap_length=80):
    wrapped_text = textwrap.fill(text, wrap_length)
    return wrapped_text

# PDF Processing
def process_pdf(uploaded_file):
    doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
    pages_and_texts = []
    
    with st.spinner("Processing PDF..."):
        for page_number, page in tqdm(enumerate(doc)):
            text = page.get_text()
            text = text_formatter(text)
            pages_and_texts.append({
                "page_number": page_number,
                "text": text,
                "char_count": len(text),
                "word_count": len(text.split(" ")),
                "token_count": len(text)/4
            })
    
    with st.spinner("Chunking text..."):
        nlp = English()
        nlp.add_pipe("sentencizer")
        for item in tqdm(pages_and_texts):
            item['sentences'] = [str(s) for s in nlp(item["text"]).sents]
            item["sentence_chunks"] = split_list(item["sentences"])
    
    pages_and_chunks = []
    for item in tqdm(pages_and_texts):
        for sentence_chunk in item["sentence_chunks"]:
            chunk_dict = {
                "page_number": item["page_number"],
                "sentence_chunk": " ".join(sentence_chunk).replace("  ", " ").strip()
            }
            chunk_dict["chunk_token_count"] = len(chunk_dict["sentence_chunk"])/4
            pages_and_chunks.append(chunk_dict)
    
    return [c for c in pages_and_chunks if c["chunk_token_count"] > MIN_TOKEN_LENGTH]

# Model Loading
@st.cache_resource
def load_models():
    embedding_model = SentenceTransformer(MODEL_NAME, device=DEVICE)
    
    # LLM Model setup
    quantization_config = BitsAndBytesConfig(load_in_4bit=True)
    tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID)
    llm_model = AutoModelForCausalLM.from_pretrained(
        LLM_MODEL_ID,
        quantization_config=quantization_config,
        device_map="auto"
    )
    return embedding_model, tokenizer, llm_model

# Streamlit UI
st.title("PDF Knowledge Assistant 📚")
st.markdown("Upload a PDF document and ask questions about its content")

# Sidebar for PDF Upload
with st.sidebar:
    st.header("Document Setup")
    uploaded_file = st.file_uploader("Upload PDF", type=["pdf"])
    if uploaded_file:
        if not st.session_state.processed:
            st.session_state.pages_and_chunks = process_pdf(uploaded_file)
            embedding_model, _, _ = load_models()
            
            with st.spinner("Generating embeddings..."):
                texts = [c["sentence_chunk"] for c in st.session_state.pages_and_chunks]
                st.session_state.embeddings = torch.tensor(
                    embedding_model.encode(texts, convert_to_tensor=True),
                    device=DEVICE
                )
            st.session_state.processed = True

# Main Q&A Interface
if st.session_state.processed:
    query = st.text_input("Enter your question about the document:")
    
    if query:
        embedding_model, tokenizer, llm_model = load_models()
        
        # Retrieve relevant chunks
        query_embedding = embedding_model.encode(query, convert_to_tensor=True)
        scores = util.dot_score(query_embedding, st.session_state.embeddings)[0]
        top_results = torch.topk(scores, k=NUM_RESULTS)
        
        # Display relevant passages
        st.subheader("Most Relevant Passages:")
        for score, idx in zip(top_results[0], top_results[1]):
            with st.expander(f"Relevance: {score:.2f}"):
                st.write(st.session_state.pages_and_chunks[idx]["sentence_chunk"])
                st.caption(f"Page: {st.session_state.pages_and_chunks[idx]['page_number']+1}")
        
        # Generate answer
        with st.spinner("Generating answer..."):
            context = [st.session_state.pages_and_chunks[i] for i in top_results[1]]
            prompt = f"""
            Answer this question: {query}
            Using information from these passages:
            {" ".join([c['sentence_chunk'] for c in context])}
            Keep answers technical and specific to the document content.
            """
            
            inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
            outputs = llm_model.generate(**inputs, max_new_tokens=500)
            answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            st.subheader("Generated Answer:")
            st.write(print_wrapped(answer.split("Answer:")[-1].strip()))
else:
    st.info("Please upload a PDF document to get started")