import gradio as gr
from pypdf import PdfReader
from rag_pipeline import RAGPipeline
import warnings
import asyncio
import logging
import os

warnings.filterwarnings("ignore", category=FutureWarning)
logging.getLogger("transformers").setLevel(logging.ERROR)
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Fix Python 3.13 asyncio cleanup bug
import sys
if sys.version_info >= (3, 13):
    import asyncio.base_events
    _original_del = asyncio.base_events.BaseEventLoop.__del__
    def _safe_del(self):
        try:
            _original_del(self)
        except Exception:
            pass
    asyncio.base_events.BaseEventLoop.__del__ = _safe_del
rag = RAGPipeline()

def process_file(file):
    if file is None:
        return "No file uploaded."
    
    # Support PDF and TXT
    if file.name.endswith(".pdf"):
        reader = PdfReader(file.name)
        text = "\n".join(page.extract_text() for page in reader.pages)
    else:
        with open(file.name, "r", encoding="utf-8") as f:
            text = f.read()

    count = rag.ingest(text)
    return f"✅ Indexed {count} chunks from your document. You can now ask questions!"

def chat(query, history):
    if not query.strip():
        return history
    answer = rag.answer(query)
    history.append((query, answer))
    return history

with gr.Blocks(title="📚 RAG App — No API Key") as demo:
    gr.Markdown("# 📚 RAG Q&A — Powered by Open-Source Models\nUpload a document, then ask questions about it.")

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(label="Upload PDF or TXT", file_types=[".pdf", ".txt"])
            upload_btn = gr.Button("📥 Process Document", variant="primary")
            status = gr.Textbox(label="Status", interactive=False)
        
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="Chat with your document")
            query_box = gr.Textbox(placeholder="Ask a question about your document...", label="Your question")
            ask_btn = gr.Button("Ask", variant="primary")

    upload_btn.click(process_file, inputs=file_input, outputs=status)
    ask_btn.click(chat, inputs=[query_box, chatbot], outputs=chatbot)
    query_box.submit(chat, inputs=[query_box, chatbot], outputs=chatbot)

demo.launch(
    server_name="0.0.0.0", 
    server_port=7860,
    ssr_mode=False,
    show_error=True,
    quiet=False)