"""
VSS Demo - Minimal Gradio App
"""

import asyncio
import base64
from pathlib import Path
from typing import Dict, Tuple

import gradio as gr
import httpx

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

client = httpx.AsyncClient(timeout=120.0)


async def api_query(text: str) -> Dict:
    r = await client.post(f"{VSS_API_URL}/query", json={"query": text})
    r.raise_for_status()
    return r.json()


async def api_query_media(text: str, image_path: str) -> Dict:
    with open(image_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    r = await client.post(f"{VSS_API_URL}/query/media", json={"query": text, "media_base64": b64})
    r.raise_for_status()
    return r.json()


def format_trace(result: Dict) -> str:
    lines = ["=" * 40, "TRACE", "=" * 40]
    meta = result.get("metadata", {})
    source = meta.get("response_source", "?")
    lines.append(f"Flow: classify -> {source} -> compose")
    lines.append("-" * 40)
    for i, step in enumerate(result.get("trace", []), 1):
        lines.append(f"{i}. {step}")
    lines.append("-" * 40)
    lines.append(f"Intent: {meta.get('intent')} | Source: {source}")
    if meta.get("sql_query"):
        lines.append(f"SQL: {meta['sql_query'][:60]}...")
    return "\n".join(lines)


def run_query(query: str, image_path: str) -> Tuple[str, str, str]:
    if not query.strip():
        return "Enter a query", "", ""
    
    async def _do():
        if image_path:
            return await api_query_media(query, image_path)
        return await api_query(query)
    
    try:
        result = asyncio.run(_do())
        resp = result.get("response", "No response")
        trace = format_trace(result)
        status = f"{result.get('status','?').upper()} | {result.get('processing_time_ms',0)}ms"
        return resp, trace, status
    except Exception as e:
        return f"Error: {e}", "", "ERROR"


def get_samples():
    samples = [""]
    for d in ["vehicle_detection", "people_detection", "fire_detection"]:
        p = DATA_DIR / d
        if p.exists():
            samples.extend([str(f) for f in list(p.glob("*.png"))[:2] + list(p.glob("*.jpg"))[:2]])
    return samples


# Build UI
with gr.Blocks(title="VSS Demo") as demo:
    gr.Markdown("# VSS Demo")
    
    with gr.Row():
        with gr.Column():
            query_box = gr.Textbox(label="Query", placeholder="Enter question...", lines=2)
            sample_dd = gr.Dropdown(label="Sample Image", choices=get_samples(), value="")
            image_box = gr.Image(label="Or Upload", type="filepath")
            run_btn = gr.Button("Run Query", variant="primary")
        
        with gr.Column():
            status_box = gr.Textbox(label="Status")
            output_box = gr.Textbox(label="Response", lines=6)
            trace_box = gr.Textbox(label="Trace", lines=12)
    
    def on_run(q, sample, uploaded):
        img = uploaded if uploaded else (sample if sample else None)
        return run_query(q, img)
    
    run_btn.click(on_run, inputs=[query_box, sample_dd, image_box], outputs=[output_box, trace_box, status_box])

demo.launch(server_name="0.0.0.0", server_port=7860)
