"""
VSS Demo - Streamlit Version
"""
import streamlit as st
import requests
import base64
from pathlib import Path

st.set_page_config(page_title="VSS Demo", layout="wide")
st.title("VSS - Video Search & Summarization")

API = "http://localhost:8080/api/v1"
DATA_DIR = Path("/workspace/AI_Agents/VLM_Agent/VLM_8xH100_export/data/frames")

# Get sample images
def get_samples():
    samples = {"None": None}
    for d in ["vehicle_detection", "people_detection", "fire_detection"]:
        p = DATA_DIR / d
        if p.exists():
            for f in list(p.glob("*.png"))[:2] + list(p.glob("*.jpg"))[:2]:
                samples[f.name] = str(f)
    return samples

col1, col2 = st.columns(2)

with col1:
    st.subheader("Query")
    query = st.text_area("Question", placeholder="Enter your question...")
    
    st.write("**Quick Examples:**")
    if st.button("SQL: How many cars?"): 
        query = "How many unique cars were detected?"
    if st.button("SQL: Vehicle breakdown"):
        query = "Show breakdown of vehicle types"
    if st.button("VLM: Describe image"):
        query = "Describe what you see in this image"
    
    st.subheader("Image (optional)")
    samples = get_samples()
    selected = st.selectbox("Sample Image", list(samples.keys()))
    uploaded = st.file_uploader("Or Upload", type=["png", "jpg", "jpeg"])
    
    run = st.button("Run Query", type="primary")

with col2:
    st.subheader("Response")
    
    if run and query:
        with st.spinner("Processing..."):
            try:
                img_path = None
                if uploaded:
                    img_path = f"/tmp/{uploaded.name}"
                    with open(img_path, "wb") as f:
                        f.write(uploaded.read())
                elif selected != "None":
                    img_path = samples[selected]
                
                if img_path:
                    with open(img_path, "rb") as f:
                        b64 = base64.b64encode(f.read()).decode()
                    r = requests.post(f"{API}/query/media", json={"query": query, "media_base64": b64})
                else:
                    r = requests.post(f"{API}/query", json={"query": query})
                
                r.raise_for_status()
                result = r.json()
                
                st.success(f"Status: {result.get('status')} | Time: {result.get('processing_time_ms')}ms")
                st.write("**Answer:**")
                st.write(result.get("response", "No response"))
                
                st.write("**Trace:**")
                meta = result.get("metadata", {})
                st.code(f"Intent: {meta.get('intent')} | Source: {meta.get('response_source')}")
                
                with st.expander("Full Trace"):
                    for i, step in enumerate(result.get("trace", []), 1):
                        st.text(f"{i}. {step}")
                
                if meta.get("sql_query"):
                    with st.expander("SQL Query"):
                        st.code(meta["sql_query"], language="sql")
                        
            except Exception as e:
                st.error(f"Error: {e}")
