import streamlit as st
import requests
import tempfile
import os

# ==============================
# API Configuration
# ==============================
API_BASE_URL = os.getenv("API_BASE_URL", "https://sameer-handsome173-multi-modal-rag.hf.space")

ENDPOINTS = {
    "ingestion": {
        "ingest_pdf": f"{API_BASE_URL}/ingest",
        "ingestion_stats": f"{API_BASE_URL}/ingest/stats",
    },
    "query": {
        "query": f"{API_BASE_URL}/query",
        "query_stats": f"{API_BASE_URL}/query/stats",
    },
    "docs": {
        "swagger": f"{API_BASE_URL}/docs",
        "redoc": f"{API_BASE_URL}/redoc",
    }
}

TIMEOUT = 150  # seconds

# ==============================
# Helper Request Function
# ==============================
def make_request(method, url, **kwargs):
    try:
        response = requests.request(method, url, timeout=TIMEOUT, **kwargs)
        response.raise_for_status()
        return response
    except requests.Timeout:
        st.error("⏳ Request timed out! The server took too long to respond.")
    except requests.RequestException as e:
        st.error(f"⚠️ Request error: {e}")
    return None

# ==============================
# Streamlit Configuration
# ==============================
st.set_page_config(page_title="Multi-Modal RAG", page_icon="🤖", layout="wide")

# ==============================
# Black-Green Theme
# ==============================
st.markdown("""
    <style>
        body {background-color: #000000;}
        .stApp {background-color: #000000; color: #00FF7F;}
        h1, h2, h3, h4 {color: #00FF7F;}
        .stButton>button {
            background-color: #00FF7F;
            color: black;
            border-radius: 10px;
            font-weight: bold;
        }
        .stButton>button:hover {
            background-color: #00cc66;
            color: white;
        }
        textarea, input {
            background-color: #1a1a1a !important;
            color: #00FF7F !important;
            border-radius: 10px !important;
        }
    </style>
""", unsafe_allow_html=True)

st.title("🤖 Multi-Modal RAG System")
st.subheader("Upload, Ingest, and Query Seamlessly")

# ==============================
# Session State
# ==============================
if "ingested" not in st.session_state:
    st.session_state.ingested = False

# ==============================
# PDF Upload and Ingestion
# ==============================
uploaded_file = st.file_uploader("📂 Upload your PDF", type=["pdf"])

if uploaded_file:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
        temp_file.write(uploaded_file.read())
        temp_path = temp_file.name

    if st.button("🚀 Ingest PDF"):
        with st.spinner("Uploading and ingesting your PDF... ⏳"):
            with open(temp_path, "rb") as f:
                files = {"file": (uploaded_file.name, f, "application/pdf")}
                response = make_request("POST", ENDPOINTS["ingestion"]["ingest_pdf"], files=files)

            if response and response.status_code == 200:
                st.success("✅ PDF ingested successfully!")
                st.session_state.ingested = True
            elif response:
                st.error(f"❌ Ingestion failed: {response.text}")

# ==============================
# Query Section
# ==============================
st.markdown("---")
st.subheader("💬 Ask a Question from the PDF")

user_query = st.text_area("Enter your query here...")

if st.button("🔍 Ask Query"):
    if not st.session_state.ingested:
        st.warning("⚠️ Please ingest a PDF before querying.")
    elif user_query.strip() == "":
        st.warning("⚠️ Please enter a question.")
    else:
        with st.spinner("Querying the model... 🤔"):
            response = make_request("POST", ENDPOINTS["query"]["query"], params={"question": user_query, "k": 5})

            if response and response.status_code == 200:
                result = response.json()
                answer = result.get("answer", "No answer found.")
                st.markdown("### 🧠 Answer:")
                st.markdown(f"**{answer}**")
            elif response:
                st.error(f"❌ Query failed: {response.text}")