import gradio as gr
import pandas as pd
import os
import time
import json
from src.vector_db import UnifiedQdrant
from src.router import LearnedRouter
from src.comparison import ComparisonEngine
from config import COLLECTION_NAME, NUM_CLUSTERS, FRESHNESS_SHARD_ID, MRL_DIMS

# --- Initialization ---
print("Initializing dashVectorspace App...")

# 1. Initialize DB
db = UnifiedQdrant(
    collection_name=COLLECTION_NAME,
    vector_size=384,
    num_clusters=NUM_CLUSTERS,
    freshness_shard_id=FRESHNESS_SHARD_ID
)
db.initialize()

# 2. Initialize Router
ROUTER_PATH = "models/router_v1.pkl"
if os.path.exists(ROUTER_PATH):
    router = LearnedRouter.load(ROUTER_PATH)
else:
    print("WARNING: Router model not found. Creating a DUMMY router for demo UI.")
    router = LearnedRouter(model_type="lightgbm", n_clusters=NUM_CLUSTERS, mrl_dims=MRL_DIMS)
    router.predict = lambda x: (0, 0.99)

# 3. Initialize Engine
engine = ComparisonEngine(db, router, embedding_model_name="minilm")

# --- UI Logic ---
def run_comparison(query):
    if not query:
        return "Please enter a query.", None, None, None, None

    res_direct = engine.direct_search(query)
    res_xvector = engine.xvector_search(query)
    
    def format_results(res_dict):
        points = res_dict["results"]
        html = "<div style='display: flex; flex-direction: column; gap: 10px;'>"
        for p in points:
            payload = p.payload
            text = payload.get("text", "No text") if payload else "No text"
            score = p.score
            # Card style for results
            html += f"""
            <div style="padding: 10px; border-radius: 8px; background: rgba(255,255,255,0.05); border: 1px solid rgba(255,255,255,0.1);">
                <div style="font-size: 0.8em; color: #aaa; margin-bottom: 4px;">Score: {score:.4f}</div>
                <div style="font-size: 0.95em;">{text[:200]}...</div>
            </div>
            """
        html += "</div>"
        return html

    out_direct = format_results(res_direct)
    out_xvector = format_results(res_xvector)
    
    # Metrics for JSON
    metrics_data = {
        "Brute Force": {
            "Latency": f"{res_direct['latency_ms']:.2f} ms",
            "Shards Searched": res_direct['shards_searched']
        },
        "xVector": {
            "Latency": f"{res_xvector['latency_ms']:.2f} ms",
            "Shards Searched": res_xvector['shards_searched'],
            "Mode": res_xvector['mode']
        }
    }
    
    # Savings
    savings = (1 - (res_xvector["shards_searched"] / res_direct["shards_searched"])) * 100
    savings_html = f"""
    <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 12px; color: white;">
        <div style="font-size: 1.2em; opacity: 0.9;">Compute Savings</div>
        <div style="font-size: 3em; font-weight: bold;">{savings:.1f}%</div>
        <div style="font-size: 0.9em; opacity: 0.8;">{res_xvector['shards_searched']} shards vs {res_direct['shards_searched']}</div>
    </div>
    """
    
    # Telemetry
    telemetry_data = {
        "Router Confidence": f"{res_xvector.get('confidence', 0):.4f}",
        "Target Cluster": int(res_xvector.get('target_cluster', -1)),
        "Search Mode": res_xvector['mode']
    }
    
    return out_direct, out_xvector, metrics_data, savings_html, telemetry_data

# --- Custom CSS ---
custom_css = """
body { background-color: #0b0f19; color: #e0e0e0; }
.gradio-container { font-family: 'Inter', sans-serif; }
h1 { background: -webkit-linear-gradient(45deg, #667eea, #764ba2); -webkit-background-clip: text; -webkit-text-fill-color: transparent; }
.result-box { border: 1px solid #333; border-radius: 8px; padding: 10px; }
"""

# --- Gradio Layout ---
with gr.Blocks(title="dashVectorspace", theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate"), css=custom_css) as demo:
    
    with gr.Column(elem_id="main-container"):
        # Header
        gr.HTML("""
        <div style="text-align: center; margin-bottom: 30px;">
            <h1 style="font-size: 3em; margin-bottom: 10px;">🚀 dashVectorspace</h1>
            <p style="font-size: 1.2em; color: #888;">Production-Grade Learned Hybrid Retrieval Engine</p>
        </div>
        """)
        
        # Input Section
        with gr.Row(variant="panel"):
            with gr.Column(scale=4):
                query_input = gr.Textbox(
                    label="Search Query", 
                    placeholder="Ask a complex question (e.g., 'How does AI impact healthcare efficiency?')", 
                    lines=1,
                    show_label=False,
                    container=False,
                    scale=4
                )
            with gr.Column(scale=1):
                submit_btn = gr.Button("🔍 Search", variant="primary", scale=1)

        # Results Section
        with gr.Row():
            # Left: Brute Force
            with gr.Column():
                gr.Markdown("### 🐢 Brute Force (Baseline)")
                out_baseline = gr.HTML(label="Results")
            
            # Right: xVector
            with gr.Column():
                gr.Markdown("### ⚡ xVector (Optimized)")
                out_optimized = gr.HTML(label="Results")

        gr.Markdown("---")

        # Metrics Section
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📊 Performance Metrics")
                metrics_display = gr.JSON(label="Detailed Metrics")
            
            with gr.Column(scale=1):
                gr.Markdown("### 💰 Efficiency")
                savings_display = gr.HTML()
        
        # Telemetry (Accordion)
        with gr.Accordion("🛠️ System Telemetry (Debug Info)", open=False):
            telemetry_display = gr.JSON(label="Router Decisions")

    # Event Listener
    submit_btn.click(
        run_comparison, 
        inputs=[query_input], 
        outputs=[out_baseline, out_optimized, metrics_display, savings_display, telemetry_display]
    )
    
    # Allow Enter key to submit
    query_input.submit(
        run_comparison, 
        inputs=[query_input], 
        outputs=[out_baseline, out_optimized, metrics_display, savings_display, telemetry_display]
    )

if __name__ == "__main__":
    demo.launch()