import gradio as gr
import torch
import faiss
import numpy as np
from PyPDF2 import PdfReader
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM

# ---------------- Step 1: Load Open-Source Models ----------------
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

# Open-Source Small Language Model (SLM)
model_name = "microsoft/phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# ---------------- Step 2: Extract & Process PDF ----------------
def extract_text_from_pdf(pdf_file_path):
    """Extract text from a PDF file, handling encryption if necessary."""
    reader = PdfReader(pdf_file_path)
    
    if reader.is_encrypted:
        try:
            reader.decrypt("")  # Attempt to decrypt without a password
        except:
            return "❌ Error: This PDF is encrypted and requires a password."
    
    text = " ".join([page.extract_text() for page in reader.pages if page.extract_text()])
    return text


def process_pdf(pdf_file):
    """Extracts text and creates embeddings for retrieval."""
    raw_text = extract_text_from_pdf(pdf_file)
    chunk_size = 500
    chunks = [raw_text[i:i + chunk_size] for i in range(0, len(raw_text), chunk_size)]
    
    tokenized_chunks = [chunk.split() for chunk in chunks]
    bm25 = BM25Okapi(tokenized_chunks)
    
    embeddings = np.array([embedding_model.encode(chunk) for chunk in chunks])
    dimension = embeddings.shape[1]
    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)
    
    return chunks, bm25, index

# ---------------- Step 3: Retrieve Top Document ----------------
def retrieve_top_doc(query, chunks, bm25, index):
    """Retrieves the single most relevant document for a query."""
    tokenized_query = query.split()
    bm25_scores = bm25.get_scores(tokenized_query)
    top_bm25_index = np.argmax(bm25_scores)

    query_embedding = embedding_model.encode(query).reshape(1, -1)
    distances, top_faiss_index = index.search(query_embedding, 1)
    
    best_index = top_bm25_index if bm25_scores[top_bm25_index] > bm25_scores[top_faiss_index[0][0]] else top_faiss_index[0][0]
    confidence_score = round(1 / (1 + distances[0][0]), 2)
    
    return chunks[best_index], confidence_score

# ---------------- Step 4: Out-of-Context Rail Guard ----------------
def is_relevant_query(confidence_score, threshold=0.2):
    """Determines if the query is relevant based on confidence score."""
    return confidence_score > threshold

def rail_guard(query):
    """Ensures that the query is within the expected domain (financial reports)."""
    finance_keywords = ["revenue", "profit", "loss", "earnings", "financial", "income", "assets", "liabilities", "balance sheet", "cash flow"]
    query_lower = query.lower()

    # If any financial keyword is present, consider the query valid
    if any(keyword in query_lower for keyword in finance_keywords):
        return True

    return False
    
# ---------------- Step 5: LLM Response Generation ----------------
def generate_response(context, query):
    """Generates a response using the LLM model."""
    prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
    
    tokenizer.pad_token = tokenizer.eos_token
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512, padding=True)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=150,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )
    
    return tokenizer.decode(output[0], skip_special_tokens=True)

# ---------------- Step 5: Gradio Interface ----------------
chunks, bm25, index = None, None, None

def handle_pdf_upload(pdf_file):
    """Processes uploaded PDF and creates retrieval index."""
    global chunks, bm25, index
    chunks, bm25, index = process_pdf(pdf_file)
    return "✅ PDF processed successfully! You can now ask financial questions."

def handle_query(query):
    """Handles user queries and generates responses."""
    if chunks is None or bm25 is None or index is None:
        return "❌ Please upload and process a PDF first.", ""

    # 🛑 Step 1: Check if the query is even financial-related
    if not rail_guard(query):
        return (
            "🚫 I'm sorry, but I can only answer questions related to financial reports. "
            "Please ask about revenue, cash flow, balance sheets, or similar topics."
        ), "🔹 Query rejected by rail guard."

    # ✅ Step 2: Retrieve a relevant document
    retrieved_doc, confidence_score = retrieve_top_doc(query, chunks, bm25, index)

    # 🛑 Step 3: If the query is out of context based on confidence, return a fallback response
    if not is_relevant_query(confidence_score):
        return (
            "🤖 I couldn't find relevant financial information in the document. "
            "Please try rephrasing your question or asking something related to financial data."
        ), "🔹 No relevant document found."

    # ✅ Step 4: Generate a response using the AI model
    response = generate_response(retrieved_doc, query)
    retrieved_doc_snippet = f"Confidence: {confidence_score}\n🔹 {retrieved_doc[:500]}..."
    
    return response, retrieved_doc_snippet

with gr.Blocks() as interface:
    gr.Markdown("## 📊 Financial Q&A - RAG on Annual Reports")
    
    with gr.Row():
        pdf_input = gr.File(label="Upload Financial Report (PDF)", type="filepath")
        process_button = gr.Button("Process PDF")
    
    pdf_status = gr.Textbox(label="PDF Status", interactive=False)
    process_button.click(handle_pdf_upload, inputs=[pdf_input], outputs=[pdf_status])
    
    query_input = gr.Textbox(label="Enter your financial question:")
    query_button = gr.Button("Get Answer")
    
    response_output = gr.Textbox(label="Answer", interactive=False)
    retrieved_docs_output = gr.Textbox(label="Retrieved Document & Confidence Score", interactive=False)
    
    query_button.click(handle_query, inputs=[query_input], outputs=[response_output, retrieved_docs_output])

interface.launch()