import streamlit as st
import pandas as pd
import json
import os
from PIL import Image

# --- CONFIG ---
st.set_page_config(layout="wide", page_title="Semantic-Drive Explorer", page_icon="🚗")
DATA_FILE = "hf_demo_pack/demo_data.jsonl"

# --- SIDEBAR ---
st.sidebar.title("🚗 Semantic-Drive")
st.sidebar.markdown("**Mining Long-Tail Edge Cases with Neuro-Symbolic VLMs**")
st.sidebar.info("This demo showcases scenarios retrieved from NuScenes using a local Llama-3/Qwen pipeline running on consumer hardware.")

# --- LOAD DATA ---
@st.cache_data
def load_data():
    if not os.path.exists(DATA_FILE):
        return pd.DataFrame()
    data = []
    with open(DATA_FILE, 'r') as f:
        for line in f:
            data.append(json.loads(line))
    return pd.DataFrame(data)

df = load_data()

if df.empty:
    st.error("No data found. Please upload demo_data.jsonl")
    st.stop()

# --- FILTERS ---
st.sidebar.header("🔍 Search Filters")

# 1. Filter by Tags
all_tags = set()
for tags in df['wod_e2e_tags']:
    all_tags.update(tags)
    
selected_tags = st.sidebar.multiselect("WOD-E2E Tags", sorted(list(all_tags)))

# 2. Filter by Risk
min_risk = st.sidebar.slider("Minimum Risk Score", 0, 10, 0)

# Apply Filters
filtered_df = df.copy()
if selected_tags:
    filtered_df = filtered_df[filtered_df['wod_e2e_tags'].apply(lambda x: any(t in x for t in selected_tags))]
filtered_df = filtered_df[filtered_df.apply(lambda x: x['scenario_criticality']['risk_score'] >= min_risk, axis=1)]

st.sidebar.markdown(f"**Found:** {len(filtered_df)} / {len(df)} scenarios")

# --- MAIN FEED ---
st.title("Scenario Database")

# Display as a feed of cards
for idx, row in filtered_df.iterrows():
    with st.container():
        c1, c2 = st.columns([1, 2])
        
        with c1:
            # Load Image
            raw_path = row.get("web_image_path", "")
            # Fix path if it doesn't start with the folder name
            if not raw_path.startswith("hf_demo_pack"):
                img_path = os.path.join("hf_demo_pack", raw_path)
            else:
                img_path = raw_path
                
            if os.path.exists(img_path):
                img = Image.open(img_path)
                st.image(img, caption=f"Token: {row['token'][:8]}...")
            else:
                st.warning("Image not available in demo pack")

        with c2:
            # Header
            tags = row['wod_e2e_tags']
            risk = row['scenario_criticality']['risk_score']
            
            # Badge Logic
            risk_color = "red" if risk >= 7 else "orange" if risk >= 4 else "green"
            st.markdown(f"### :{risk_color}[Risk {risk}/10]  {' '.join([f'`{t}`' for t in tags])}")
            
            # Description
            st.info(row['description'])
            
            # Details Expander
            with st.expander("🧬 View Scenario DNA (JSON)"):
                st.json({
                    "ODD": row['odd_attributes'],
                    "Topology": row['road_topology'],
                    "Agents": row['key_interacting_agents'],
                    "Reasoning": row.get('judge_log', [])
                })
        
        st.divider()