import warnings
warnings.filterwarnings("ignore")

import numpy as np
import torch
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
import faiss
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import gradio as gr
from pdf2image import convert_from_path
import pytesseract
import os

# Global variables
vector_store = None
embedding_model = None
tokenizer = None
model = None
document_category = None
document_text = None

# -------------------- Text splitter --------------------
def split_text_into_chunks(text, chunk_size=500, chunk_overlap=50):
    text = text.replace('\n', ' ')
    sentences = text.split('. ')
    chunks = []
    current = ""
    for sent in sentences:
        sent = sent.strip() + '. '
        if len(current) + len(sent) <= chunk_size:
            current += sent
        else:
            if current:
                chunks.append(current.strip())
            current = sent
    if current:
        chunks.append(current.strip())
    if not chunks:
        chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size - chunk_overlap)]
    return chunks

# -------------------- Vector store --------------------
class VectorStore:
    def __init__(self, embed_model):
        self.embed_model = embed_model
        self.index = None
        self.chunks = []

    def add_chunks(self, chunks):
        self.chunks = chunks
        embeddings = self.embed_model.encode(chunks, show_progress_bar=True)
        dim = embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dim)
        self.index.add(embeddings.astype(np.float32))

    def search(self, query, k=3):
        q_emb = self.embed_model.encode([query])
        dist, idx = self.index.search(q_emb.astype(np.float32), k)
        return [self.chunks[i] for i in idx[0]]

# -------------------- Language model --------------------
def load_language_model():
    model_name = "google/flan-t5-small"
    print("Loading Q&A model...")
    tok = AutoTokenizer.from_pretrained(model_name)
    mdl = AutoModelForSeq2SeqLM.from_pretrained(
        model_name,
        device_map="cpu",
        torch_dtype=torch.float32
    )
    return tok, mdl

def generate_answer(prompt, tokenizer, model, max_new_tokens=150):
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.1,
            do_sample=False
        )
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer.strip()

# -------------------- Document classification --------------------
def load_classifier():
    print("Loading document classifier...")
    classifier = pipeline(
        "zero-shot-classification",
        model="valhalla/distilbart-mnli-12-3",
        device=-1
    )
    return classifier

def classify_document(text, classifier):
    candidate_labels = [
        "invoice", "research paper", "legal contract", "financial report",
        "business proposal", "technical documentation", "resume", "general document"
    ]
    sample_text = text[:3000]
    result = classifier(sample_text, candidate_labels)
    top_label = result['labels'][0]
    top_score = result['scores'][0]
    return top_label, top_score, result

# -------------------- PDF extraction with OCR --------------------
def extract_text_from_pdf(pdf_path):
    # Try pypdf first
    try:
        reader = PdfReader(pdf_path)
        text = ""
        for page in reader.pages:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
        if text.strip():
            return text
    except Exception as e:
        print("pypdf failed:", e)

    # If no text, give a clear message (no OCR attempt)
    return "❌ This PDF is scanned or does not contain selectable text. Please upload a text‑based PDF."

    # Fallback: OCR
    print("No selectable text found. Running OCR...")
    images = convert_from_path(pdf_path, dpi=300)
    text = ""
    for img in images:
        page_text = pytesseract.image_to_string(img)
        text += page_text + "\n"
    return text

# -------------------- QA function --------------------
def answer_question(query, vector_store, tokenizer, model, category):
    context_chunks = vector_store.search(query, k=3)
    context = "\n\n".join(context_chunks)
    prompt = f"""You are analyzing a {category} document. Answer the question based only on the context below.

Document Type: {category}

Context:
{context}

Question: {query}
Answer:"""
    return generate_answer(prompt, tokenizer, model)

# -------------------- Lazy model loading --------------------
def get_models():
    global embedding_model, tokenizer, model
    if embedding_model is None:
        print("Loading embedding model...")
        embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
    if tokenizer is None or model is None:
        tokenizer, model = load_language_model()
    return embedding_model, tokenizer, model

# -------------------- Gradio handlers --------------------
def process_pdf(file):
    global vector_store, document_category, document_text
    try:
        embed_model, _, _ = get_models()
        text = extract_text_from_pdf(file.name)
        if not text.strip():
            return "❌ No text could be extracted. The PDF may be image-based and OCR failed."
        document_text = text
        print("Classifying document...")
        classifier = load_classifier()
        category, confidence, _ = classify_document(text, classifier)
        document_category = category
        print(f"Document classified as: {category} (confidence: {confidence:.2f})")
        chunks = split_text_into_chunks(text)
        print(f"Created {len(chunks)} chunks.")
        vs = VectorStore(embed_model)
        vs.add_chunks(chunks)
        vector_store = vs
        return f"✅ PDF processed!\n📄 Document Type: {category} (confidence: {confidence:.1%})\n📑 {len(chunks)} chunks created.\n\nYou can now ask questions about this {category}."
    except Exception as e:
        import traceback
        traceback.print_exc()
        return f"❌ Error: {str(e)}"

def answer(query):
    global vector_store, document_category
    if vector_store is None:
        return "Please upload and process a PDF first."
    _, tokenizer, model = get_models()
    return answer_question(query, vector_store, tokenizer, model, document_category or "general document")

# -------------------- Gradio interface --------------------
with gr.Blocks(title="Intelligent PDF Q&A Bot with OCR", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📄 Intelligent PDF Q&A Bot with OCR & Document Classification")
    gr.Markdown("Upload any PDF (text-based or scanned) – the system will automatically detect the document type and answer your questions.")
    with gr.Row():
        pdf_input = gr.File(label="📁 Upload PDF", file_types=[".pdf"])
        process_btn = gr.Button("Process PDF")
        status = gr.Textbox(label="Status", interactive=False, lines=3)
    process_btn.click(process_pdf, inputs=pdf_input, outputs=status)
    gr.Markdown("## Ask a Question")
    gr.Markdown("The AI uses the document's content and its detected category to provide better answers.")
    question = gr.Textbox(label="Your Question", placeholder="e.g., What is the total amount? (for invoices) or What is the main finding? (for research papers)")
    answer_output = gr.Textbox(label="Answer", interactive=False, lines=5)
    ask_btn = gr.Button("Ask")
    ask_btn.click(answer, inputs=question, outputs=answer_output)
    question.submit(answer, inputs=question, outputs=answer_output)
    gr.Markdown("---")
    gr.Markdown("### Supported Document Types")
    gr.Markdown("Invoices, Research Papers, Legal Contracts, Financial Reports, Business Proposals, Technical Documentation, Resumes, and General Documents.")

demo.launch(server_name="0.0.0.0", server_port=7860)