import gradio as gr
import requests
import uuid
import os
from datetime import datetime

PRESIGN_API_URL = os.environ.get("PRESIGN_API_URL")     
INFERENCE_API_URL = os.environ.get("INFERENCE_API_URL")

user_id = "user123"  
session_id = datetime.now().strftime("%Y%m%d_%H%M%S") + "_" + str(uuid.uuid4())[:8]

upload_success = False

def upload_pdf(file):
    global upload_success
    filename = os.path.basename(file.name)
    content_type = "application/pdf"
    response = requests.post(url=PRESIGN_API_URL, json={
        "filename": filename,
        "user_id": user_id,
        "session_id": session_id,
        "contentType": content_type
    })   
    if response.status_code != 200:
        return "❌ Failed to get presigned URL"
    data = response.json()
    presigned_url = data["url"]
    content_type = data.get("content_type", "application/pdf")

    with open(file.name, "rb") as f:
        upload_resp = requests.put(presigned_url, data=f, headers={"Content-Type": content_type})
    if upload_resp.status_code == 200:
        upload_success = True
        return f"✅ Uploaded and indexing triggered! You can now ask questions."
    else:
        return "❌ Upload failed!"

def chat_with_doc(message, history):
    if not upload_success:
        return "Please upload a PDF first.", history

    if message.lower() == "exit":
        return "🛑 Session ended. Thank you!", []

    payload = {
        "question": message,
        "user_id": user_id,
        "session_id": session_id,
        "chat_history": history
    }
    resp = requests.post(INFERENCE_API_URL, json=payload)  
    if resp.status_code != 200:
        return "❌ Error calling LLM. Try again later.", history
    answer = resp.json()["answer"]
    history.append((message, answer))
    return "", history

with gr.Blocks() as demo:
    gr.Markdown("## 📄 Upload your PDF and Ask Questions")
    
    with gr.Row():
        file_input = gr.File(label="Upload PDF", file_types=[".pdf"])
        upload_button = gr.Button("Upload")
        upload_output = gr.Textbox(label="Upload Status")

    upload_button.click(upload_pdf, inputs=file_input, outputs=upload_output)

    chatbot = gr.Chatbot(label="Document Chat")
    question = gr.Textbox(label="Ask a question (or type 'exit' to end)", placeholder="Ask something about the document...")
    ask_button = gr.Button("Send")

    ask_button.click(chat_with_doc, inputs=[question, chatbot], outputs=[question, chatbot])

demo.launch()