import os
import gradio as gr
from groq import Groq
from fpdf import FPDF

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings

# -------------------------
# GROQ CLIENT
# -------------------------
client = Groq(api_key=os.environ.get("RAG_API_KEY"))

# -------------------------
# GLOBAL STORAGE
# -------------------------
vectorstore = None
chat_history = []

# -------------------------
# LOAD DOCUMENTS
# -------------------------
def load_documents(files):
    documents = []
    for file in files:
        if file.name.endswith(".pdf"):
            loader = PyPDFLoader(file.name)
        else:
            loader = TextLoader(file.name)
        documents.extend(loader.load())
    return documents

# -------------------------
# BUILD VECTOR STORE
# -------------------------
def build_vectorstore(files):
    global vectorstore
    docs = load_documents(files)

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=100
    )
    chunks = splitter.split_documents(docs)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    vectorstore = FAISS.from_documents(chunks, embeddings)
    return "✅ Documents processed successfully."

# -------------------------
# ASK QUESTION
# -------------------------
def ask_question(question):
    global chat_history, vectorstore

    if vectorstore is None:
        return "❌ Please upload documents first.", ""

    docs = vectorstore.similarity_search(question, k=4)
    context = "\n\n".join([d.page_content for d in docs])

    prompt = f"""
You are a helpful AI assistant.

First, use the document context below to answer.
If the document does not fully answer the question,
you may add relevant general knowledge.

Document Context:
{context}

Question:
{question}
"""

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[{"role": "user", "content": prompt}],
    )

    answer = response.choices[0].message.content

    chat_history.append((question, answer))

    sources = "\n\n".join(
        [f"Source {i+1}:\n{d.page_content[:300]}..." for i, d in enumerate(docs)]
    )

    return answer, sources

# -------------------------
# EXPORT CHAT TO PDF
# -------------------------
def export_chat_pdf():
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)

    pdf.cell(0, 10, "AneesLLM - Chat Export", ln=True)
    pdf.ln(5)

    for q, a in chat_history:
        pdf.multi_cell(0, 8, f"Q: {q}")
        pdf.multi_cell(0, 8, f"A: {a}")
        pdf.ln(4)

    path = "/tmp/AneesLLM_Chat.pdf"
    pdf.output(path)
    return path

# -------------------------
# GRADIO UI
# -------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🤖 AneesLLM – RAG Based AI Assistant  
        Upload documents and ask intelligent questions with sources.
        """
    )

    with gr.Row():
        file_input = gr.File(
            file_types=[".pdf", ".txt"],
            file_count="multiple",
            label="📄 Upload Documents"
        )
        process_btn = gr.Button("📚 Process Documents")

    status = gr.Textbox(label="Status")

    process_btn.click(
        build_vectorstore,
        inputs=file_input,
        outputs=status
    )

    gr.Markdown("## 💬 Ask a Question")

    question = gr.Textbox(
        placeholder="Ask something about your documents...",
        label="Your Question"
    )

    ask_btn = gr.Button("Ask")

    answer = gr.Textbox(label="Answer", lines=6)
    sources = gr.Textbox(label="Sources Used", lines=6)

    ask_btn.click(
        ask_question,
        inputs=question,
        outputs=[answer, sources]
    )

    export_btn = gr.Button("📄 Export Chat as PDF")
    pdf_file = gr.File(label="Download PDF")

    export_btn.click(export_chat_pdf, outputs=pdf_file)

demo.launch()