import os
os.environ["OMP_NUM_THREADS"] = "1"

# Suppress all warnings BEFORE importing anything else
import warnings
warnings.filterwarnings('ignore')
warnings.simplefilter('ignore')

import asyncio
from datetime import datetime
import sys
import platform
import json

import pandas as pd
import gradio as gr
import folium
import base64
from datetime import timedelta
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
import io
from PIL import Image

from detector import detect_plate

from database import (
    save_detection,
    run_query,
    health_check,
    get_vehicles_by_state,
    get_hourly_traffic,
    get_top_plates,
    get_suspicious_vehicles,
    HF_TOKEN,
    DATABASE_URL,
    client,
    engine
)

from vehicle_map import (
    search_vehicle_route,
    default_vehicle_map,
    map_to_html
)

from ai_investigation import (
    ask_investigation_question,
    format_investigation_output
)

# =========================================================
# ASYNC FIX & EVENT LOOP MANAGEMENT
# =========================================================

# Suppress all warnings to prevent asyncio event loop cleanup messages
warnings.filterwarnings('ignore')
warnings.filterwarnings('ignore', category=ResourceWarning)
warnings.filterwarnings('ignore', category=RuntimeWarning)
warnings.filterwarnings('ignore', category=DeprecationWarning)

# Suppress asyncio event loop cleanup errors on Linux
import sys
_original_excepthook = sys.excepthook
def _suppress_asyncio_errors(exc_type, exc_value, traceback):
    """Suppress asyncio cleanup errors but show real errors"""
    if exc_type and 'Invalid file descriptor' in str(exc_value):
        return  # Silently ignore asyncio cleanup errors
    if exc_type and 'BaseEventLoop' in str(traceback):
        return  # Silently ignore event loop cleanup
    _original_excepthook(exc_type, exc_value, traceback)

sys.excepthook = _suppress_asyncio_errors

# =========================================================
# INVESTIGATION ASSISTANT WRAPPERS
# =========================================================

def ask_question_wrapper(question):
    """Unified AI Investigation Agent - handles any natural language question"""
    if not question or len(question.strip()) < 2:
        return ("Please ask a valid question", "", "", None, "")
    
    result = ask_investigation_question(question)
    if result.get("status") == "error":
        error_msg = result.get("message", "Investigation failed")
        return (f"Error: {error_msg}", "", "", None, "")
    
    return (
        result.get("answer", ""),
        "\n".join(result.get("findings", [])),
        result.get("analysis", {}),
        None,
        result.get("data_summary", "")
    )




# =========================================================
# DETECTION
# =========================================================

def detect_and_save(image):

    if image is None:

        return (
            "No image uploaded",
            {"error": "No image uploaded"}
        )

    try:

        now = datetime.now()

        date = now.strftime("%Y-%m-%d")
        time = now.strftime("%H:%M:%S")

        plate, state, vehicle_type, vehicle_conf, success = detect_plate(image)

        if success and plate:

            save_detection(
                plate,
                state,
                vehicle_type,
                vehicle_conf,
                date,
                time
            )

        result_text = f"""
Detection Success

Date: {date}
Time: {time}

Vehicle Type: {vehicle_type}
Plate: {plate}
State: {state}

Confidence: {round(vehicle_conf, 3)}
Saved: {success}
"""

        result_json = {
            "date": date,
            "time": time,
            "plate": plate,
            "state": state,
            "vehicle_type": vehicle_type,
            "confidence": round(vehicle_conf, 3),
            "saved": success
        }

        return result_text, result_json

    except Exception as e:

        return (
            f"Error: {str(e)}",
            {"error": str(e)}
        )

# =========================================================
# NLP QUERY
# =========================================================

def query_database(user_query):

    try:

        print(f"\n🔍 NLP Query: {user_query}")

        if not user_query.strip():

            print("❌ Empty query")
            return (
                "",
                pd.DataFrame(),
                {"error": "❌ Empty query - Please enter something"}
            )

        if not HF_TOKEN:
            print("❌ HF_TOKEN missing")
            return (
                "",
                pd.DataFrame(),
                {"error": "❌ HF_TOKEN not set - NLP features disabled"}
            )

        if not DATABASE_URL:
            print("❌ DATABASE_URL missing")
            return (
                "",
                pd.DataFrame(),
                {"error": "❌ DATABASE_URL not set - Database features disabled"}
            )

        print(f"Calling run_query...")
        response = run_query(user_query)

        print(f"Response: {response}")

        sql = response.get("sql", "")
        results = response.get("result", [])
        error = response.get("error", "")

        if error:
            print(f"❌ Query error: {error}")
            return (
                sql,
                pd.DataFrame(),
                {"error": f"❌ Query Error: {error}"}
            )

        if len(results) > 0:
            df = pd.DataFrame(results)
            print(f"✅ Found {len(results)} records")
        else:
            df = pd.DataFrame({
                "message": ["No results found"]
            })
            print("ℹ️ No results found")

        return (
            sql,
            df,
            {"success": f"✅ Found {len(results)} records"}
        )

    except Exception as e:

        error_msg = f"❌ Error: {str(e)}"
        print(error_msg)
        import traceback
        traceback.print_exc()
        return (
            "",
            pd.DataFrame(),
            {"error": error_msg}
        )

# =========================================================
# CHATBOT
# =========================================================

def chatbot_query(message, history):

    try:

        print(f"\n💬 Chatbot message: {message}")
        print(f"History type: {type(history)}, Length: {len(history) if history else 0}")

        if not message.strip():
            print("❌ Empty message")
            if not history:
                history = []
            # Gradio 6.14.0+: Format MUST be {"role": "user/assistant", "content": "text"}
            return history + [{"role": "user", "content": message}, {"role": "assistant", "content": "❌ Empty message"}], ""

        if not HF_TOKEN:
            print("❌ HF_TOKEN not configured")
            if not history:
                history = []
            return history + [{"role": "user", "content": message}, {"role": "assistant", "content": "❌ HF_TOKEN not configured - NLP disabled"}], ""

        if not DATABASE_URL:
            print("❌ DATABASE_URL not configured")
            if not history:
                history = []
            return history + [{"role": "user", "content": message}, {"role": "assistant", "content": "❌ DATABASE_URL not configured - Database disabled"}], ""

        print("Calling run_query...")
        response = run_query(message)
        print(f"Response: {response}")

        sql = response.get("sql", "")
        results = response.get("result", [])
        error = response.get("error", "")
        count = response.get("count", 0)

        if not history:
            history = []

        if error:
            print(f"❌ Query error: {error}")
            bot_reply = f"❌ Error: {error}"
        else:
            preview = str(results[:3]) if results else "No results"
            print(f"✅ Found {count} records")
            bot_reply = f"""✅ Query processed

**SQL Generated:**
```sql
{sql}
```

**Results:** {count} records found
"""

        # Gradio 6.14.0+: Append dict format {"role", "content"}
        history = history + [{"role": "user", "content": message}, {"role": "assistant", "content": bot_reply}]

        print(f"Returning history with {len(history)} messages")
        return history, ""

    except Exception as e:

        print(f"❌ Exception: {str(e)}")
        import traceback
        traceback.print_exc()

        if not history:
            history = []

        history = history + [[message, f"❌ Error: {str(e)}"]]

        return history, ""

# =========================================================
# ANALYTICS
# =========================================================

# ========== CHARTING FUNCTIONS ==========

def create_state_chart(state_df):
    """Create pie chart for vehicles by state"""
    try:
        if state_df.empty or len(state_df) == 0:
            return None
        
        plt.figure(figsize=(10, 6))
        
        # Get state and count columns by name
        if 'state' in state_df.columns and 'count' in state_df.columns:
            states = state_df['state'].tolist()
            counts = state_df['count'].tolist()
        else:
            states = state_df.iloc[:, 0].tolist() if len(state_df.columns) > 0 else []
            counts = state_df.iloc[:, 1].tolist() if len(state_df.columns) > 1 else []
        
        if not states or not counts:
            return None
        
        # Create pie chart
        colors = plt.cm.Set3(range(len(states)))
        plt.pie(counts, labels=states, autopct='%1.1f%%', colors=colors, startangle=90)
        plt.title('🚗 Vehicles by State Distribution', fontsize=14, fontweight='bold', pad=20)
        plt.tight_layout()
        
        # Convert to PIL Image
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
        buf.seek(0)
        img = Image.open(buf)
        img_copy = img.copy()
        plt.close()
        return img_copy
    except Exception as e:
        print(f"State chart error: {e}")
        plt.close()
        return None


def create_hourly_chart(hourly_df):
    """Create line chart for hourly traffic"""
    try:
        if hourly_df.empty or len(hourly_df) == 0:
            return None
        
        plt.figure(figsize=(12, 5))
        
        # Get hour and traffic columns by name
        if 'hour' in hourly_df.columns and 'traffic' in hourly_df.columns:
            hours = hourly_df['hour'].tolist()
            counts = hourly_df['traffic'].tolist()
        else:
            hours = hourly_df.iloc[:, 0].tolist() if len(hourly_df.columns) > 0 else []
            counts = hourly_df.iloc[:, 1].tolist() if len(hourly_df.columns) > 1 else []
        
        if not hours or not counts:
            return None
        
        # Create line chart
        plt.plot(hours, counts, marker='o', linewidth=2, markersize=8, color='#FF6B6B')
        plt.fill_between(range(len(hours)), counts, alpha=0.3, color='#FF6B6B')
        plt.xlabel('Hour of Day', fontsize=11, fontweight='bold')
        plt.ylabel('Detection Count', fontsize=11, fontweight='bold')
        plt.title('📊 Traffic by Hour', fontsize=14, fontweight='bold', pad=20)
        plt.grid(True, alpha=0.3)
        plt.xticks(range(0, len(hours), max(1, len(hours)//12)))
        plt.tight_layout()
        
        # Convert to PIL Image
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
        buf.seek(0)
        img = Image.open(buf)
        img_copy = img.copy()
        plt.close()
        return img_copy
    except Exception as e:
        print(f"Hourly chart error: {e}")
        plt.close()
        return None


def create_top_plates_chart(top_df):
    """Create horizontal bar chart for top plates"""
    try:
        if top_df.empty or len(top_df) == 0:
            return None
        
        plt.figure(figsize=(10, 6))
        
        # Get plate and detections columns (limit to top 10)
        if 'plate' in top_df.columns and 'detections' in top_df.columns:
            plates = top_df.head(10)['plate'].tolist()
            counts = top_df.head(10)['detections'].tolist()
        else:
            plates = top_df.iloc[:10, 0].tolist() if len(top_df.columns) > 0 else []
            counts = top_df.iloc[:10, 1].tolist() if len(top_df.columns) > 1 else []
        
        if not plates or not counts:
            return None
        
        # Create horizontal bar chart
        colors = plt.cm.viridis(range(len(plates)))
        bars = plt.barh(plates, counts, color=colors)
        
        # Add value labels on bars
        for i, (bar, count) in enumerate(zip(bars, counts)):
            plt.text(count + 0.1, i, str(int(count)), va='center', fontsize=9)
        
        plt.xlabel('Detection Count', fontsize=11, fontweight='bold')
        plt.title('🏆 Top Detected License Plates', fontsize=14, fontweight='bold', pad=20)
        plt.tight_layout()
        
        # Convert to PIL Image
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
        buf.seek(0)
        img = Image.open(buf)
        img_copy = img.copy()
        plt.close()
        return img_copy
    except Exception as e:
        print(f"Top plates chart error: {e}")
        plt.close()
        return None


def create_suspicious_chart(suspicious_df):
    """Create donut chart for suspicious vehicles"""
    try:
        if suspicious_df.empty or len(suspicious_df) == 0:
            return None
        
        plt.figure(figsize=(10, 6))
        
        # Get data by column name (limit to top 8)
        if 'plate' in suspicious_df.columns and 'detections' in suspicious_df.columns:
            labels = suspicious_df.head(8)['plate'].tolist()
            sizes = suspicious_df.head(8)['detections'].tolist()
        else:
            labels = suspicious_df.iloc[:8, 0].tolist() if len(suspicious_df.columns) > 0 else []
            sizes = suspicious_df.iloc[:8, 1].tolist() if len(suspicious_df.columns) > 1 else []
        
        if not labels or not sizes:
            return None
        
        # Create donut chart
        colors = plt.cm.Reds(range(len(labels)))
        wedges, texts, autotexts = plt.pie(sizes, labels=labels, autopct='%1.1f%%', 
                                           colors=colors, startangle=90, pctdistance=0.85)
        
        # Draw donut hole
        centre_circle = plt.Circle((0, 0), 0.70, fc='white')
        plt.gca().add_artist(centre_circle)
        
        plt.title('⚠️ Suspicious Vehicle Alerts', fontsize=14, fontweight='bold', pad=20)
        plt.tight_layout()
        
        # Convert to PIL Image
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
        buf.seek(0)
        img = Image.open(buf)
        img_copy = img.copy()
        plt.close()
        return img_copy
    except Exception as e:
        print(f"Suspicious chart error: {e}")
        plt.close()
        return None


def refresh_analytics():

    try:

        print("\n📊 Refreshing analytics...")

        if not DATABASE_URL or not engine:
            print("❌ Database not configured")
            err = pd.DataFrame({"error": ["Database not configured"]})
            return err, err, err, err, err, err, err, err

        print("Getting vehicles by state...")
        state_data = get_vehicles_by_state()
        state_df = pd.DataFrame(state_data) if state_data else pd.DataFrame()

        print("Getting hourly traffic...")
        hourly_data = get_hourly_traffic()
        hourly_df = pd.DataFrame(hourly_data) if hourly_data else pd.DataFrame()

        print("Getting top plates...")
        top_data = get_top_plates()
        top_df = pd.DataFrame(top_data) if top_data else pd.DataFrame()

        print("Getting suspicious vehicles...")
        suspicious_data = get_suspicious_vehicles()
        suspicious_df = pd.DataFrame(suspicious_data) if suspicious_data else pd.DataFrame()

        print(f"✅ Analytics refreshed: {len(state_df)} states, {len(hourly_df)} hours, {len(top_df)} top plates, {len(suspicious_df)} suspicious")

        # Generate charts
        print("🎨 Generating charts...")
        state_chart = create_state_chart(state_df)
        hourly_chart = create_hourly_chart(hourly_df)
        top_chart = create_top_plates_chart(top_df)
        suspicious_chart = create_suspicious_chart(suspicious_df)

        return (
            state_chart,
            hourly_chart,
            top_chart,
            suspicious_chart,
            state_df,
            hourly_df,
            top_df,
            suspicious_df
        )

    except Exception as e:

        print(f"❌ Analytics error: {str(e)}")
        import traceback
        traceback.print_exc()

        err_df = pd.DataFrame({
            "error": [str(e)]
        })

        return (
            None,
            None,
            None,
            None,
            err_df,
            err_df,
            err_df,
            err_df
        )

# =========================================================
# HEALTH
# =========================================================

status, msg = health_check()


# =========================================================
# CONFIGURATION DEBUG
# =========================================================

print("\n" + "="*60)
print("🔧 CONFIGURATION STATUS")
print("="*60)
print(f"HF_TOKEN: {'✅ SET' if HF_TOKEN else '❌ NOT SET'}")
print(f"DATABASE_URL: {'✅ SET' if DATABASE_URL else '❌ NOT SET'}")
print(f"Mistral Client: {'✅ READY' if client else '❌ FAILED'}")
print(f"Database Engine: {'✅ READY' if engine else '❌ FAILED'}")
print(f"Database Health: {msg}")
print("="*60 + "\n")

# =========================================================
# UI
# =========================================================


with gr.Blocks(
    title="ActionSync — Vehicle Intelligence Platform",
    theme=gr.themes.Soft(
        text_size="lg",
    ),
    css="""
    .status-badge {
        display: inline-flex;
        align-items: center;
        gap: 0.5em;
        padding: 0.4em 0.8em;
        border-radius: 999px;
        font-size: 0.9em;
        font-weight: 500;
    }
    .status-success {
        background-color: #16a34a;
        color: white;
        border: 1px solid #15803d;
    }
    .status-error {
        background-color: #b91c1c;
        color: white;
        border: 1px solid #991b1b;
    }
    .status-warn {
        background-color: #ca8a04;
        color: white;
        border: 1px solid #a16207;
    }
    .status-label {
        display: inline-block;
        margin-left: 0.5em;
    }
    """
) as demo:

    gr.Markdown("""
    # 🚦 ActionSync — Vehicle Intelligence Platform

    AI‑powered vehicle detection and analytics with NLP‑to‑SQL querying.
    """, elem_classes=["text-center"])

    # Status panel
    with gr.Accordion("🔧 System Status", open=True):
        with gr.Row():
            # LLM / NLP status
            with gr.Column():
                if HF_TOKEN and client:
                    gr.Markdown("""
                    <div class="status-badge status-success">
                        ✅ NLP Engine
                        <span class="status-label">Mistral LLM enabled</span>
                    </div>
                    """, elem_classes=["status-success"])
                else:
                    gr.Markdown("""
                    <div class="status-badge status-warn">
                        ⚠️ NLP Engine
                        <span class="status-label">HF_TOKEN missing</span>
                    </div>
                    """, elem_classes=["status-warn"])

            # Database status
            with gr.Column():
                if DATABASE_URL and engine:
                    # Example record count (you can replace with actual count later)
                    record_count = 20633  # ← replace with a real count query if you want
                    gr.Markdown(f"""
                    <div class="status-badge status-success">
                        ✅ Database
                        <span class="status-label">Connected — {record_count:,} records</span>
                    </div>
                    """, elem_classes=["status-success"])
                else:
                    gr.Markdown("""
                    <div class="status-badge status-error">
                        ❌ Database
                        <span class="status-label">Not configured</span>
                    </div>
                    """, elem_classes=["status-error"])

        # Setup note if anything is missing
        if not (HF_TOKEN and client) or not (DATABASE_URL and engine):
            gr.Markdown("""
            ### ⚠️ Setup Required

            To enable all features, please configure your Hugging Face and database settings:

            1. Go to **Space Settings → Repository secrets**
            2. Add:
               - `HF_TOKEN` — from [Hugging Face Tokens](https://huggingface.co/settings/tokens)
               - `DATABASE_URL` — PostgreSQL URL: `postgresql://user:pass@host:5432/db`
            3. **Restart the Space**

            After configuration, NLP queries and database features will be enabled.
            """)

    gr.Markdown(f"### {msg}")


    # =====================================================
    # TAB 1: Detection
    # =====================================================

    with gr.Tab("Detection", id="tab_detection"):
        gr.Markdown("### Vehicle Detection")

        with gr.Row():
            with gr.Column(scale=1):
                input_img = gr.Image(
                    type="numpy",
                    label="Upload Image",
                    height=360,
                )
                detect_btn = gr.Button(
                    "🔍 Detect Vehicle",
                    variant="primary",
                    size="lg",
                )

            with gr.Column(scale=1):
                output_text = gr.Textbox(
                    label="Detection Result",
                    lines=8,
                    show_label=True,
                )
                output_json = gr.JSON(
                    label="Raw JSON Output",
                )

        detect_btn.click(
            fn=detect_and_save,
            inputs=input_img,
            outputs=[
                output_text,
                output_json
            ]
        )


    # =====================================================
    # TAB 2: NLP Query
    # =====================================================

    with gr.Tab("NLP Query", id="tab_nlp"):
        gr.Markdown("### Natural Language Database Query")

        query_input = gr.Textbox(
            label="Enter your question",
            placeholder="Show vehicles from Tamil Nadu last 24 hours",
            lines=2,
        )

        with gr.Row():
            search_btn = gr.Button(
                "📝 Generate & Execute SQL",
                variant="primary",
                size="lg",
            )

        sql_output = gr.Code(
            language="sql",
            label="Generated SQL",
            lines=6,
        )

        results_output = gr.Dataframe(
            label="Query Results",
            interactive=False,
        )

        json_output = gr.JSON(
            label="Full Response",
        )

        search_btn.click(
            fn=query_database,
            inputs=query_input,
            outputs=[
                sql_output,
                results_output,
                json_output
            ]
        )


    # =====================================================
    # TAB 3: AI Investigation Agent (CHAT INTERFACE)
    # =====================================================

    with gr.Tab("🔍 AI Investigation", id="tab_investigation"):
        
        gr.Markdown("## 🤖 AI Investigation ChatBot")
        gr.Markdown("*Ask questions about vehicles, locations, patterns. Follow-ups appear naturally in chat.*")
        
        # Conversation state to track history
        conversation_state = gr.State([])
        investigation_results = gr.State({})
        
        with gr.Row():
            with gr.Column(scale=5):
                # Chat display area
                chatbot = gr.Chatbot(
                    label="💬 Investigation Conversation",
                    height=500
                )
            
            with gr.Column(scale=2):
                # Quick stats panel
                stats_display = gr.Markdown(
                    value="**📊 Investigation Stats**\n\n*Results will appear here*",
                    label="Quick Stats"
                )
        
        # Input area with follow-up suggestions
        with gr.Row():
            question_input = gr.Textbox(
                placeholder="e.g., show bikes in adyar or TN57GR8753 or midnight detections",
                label="Your Question",
                lines=2,
                scale=5
            )
            
            investigate_btn = gr.Button("🔎 Investigate", variant="primary", scale=1, size="lg")
        
        # Suggested follow-ups (appears dynamically)
        suggested_follow_ups = gr.Markdown(
            value="",
            label="💡 Suggested Questions"
        )
        
        # Detailed investigation tabs (collapsible section)
        with gr.Accordion("📋 Detailed Investigation Results", open=True):
            
            with gr.Tabs():
                
                # Data collected tab
                with gr.Tab("📊 Data Summary"):
                    data_structure = gr.Markdown(
                        value="*Investigation details will appear here*"
                    )
                
                # Key findings tab
                with gr.Tab("🚨 Key Findings & Patterns"):
                    findings_display = gr.Markdown(
                        value="*Key findings from analysis*"
                    )
                
                # Relevant data table
                with gr.Tab("📋 All Records"):
                    data_table = gr.Dataframe(
                        label="Retrieved Data",
                        interactive=False
                    )
                
                # Analysis metrics
                with gr.Tab("📈 Analysis Metrics"):
                    metrics_display = gr.Markdown(
                        value="*Metrics and statistics*"
                    )
        
        # =====================================================
        # Chat Interface Logic - FIXED MESSAGE FORMAT
        # =====================================================
        
        def investigate_fast(message, chat_history, conv_state, inv_results):
            """Fast response - append user message only"""
            if not message or len(str(message).strip()) < 2:
                return (chat_history or []), (conv_state or []), (inv_results or {})
            
            # Ensure chat_history is a list
            if not chat_history:
                history = []
            else:
                history = list(chat_history) if isinstance(chat_history, list) else []
            
            msg_str = str(message).strip()
            
            # SIMPLE: Only append user message dict, assistant will be added in next step
            user_msg_dict = {"role": "user", "content": msg_str}
            history.append(user_msg_dict)
            
            return history, (conv_state or []) + [msg_str], (inv_results or {})
        
        def investigate_complete(message):
            """Get investigation results"""
            try:
                result = ask_investigation_question(message)
                return result
            except Exception as e:
                return {
                    "status": "error",
                    "message": f"Error: {str(e)[:80]}",
                    "analysis": {},
                    "data_preview": [],
                    "total_records": 0
                }
        
        def update_chat_final(message, chat_history, inv_results):
            """Add assistant response to chat"""
            if not message or not chat_history:
                return (chat_history or []), (inv_results or {})
            
            # Ensure chat_history is a list
            history = list(chat_history) if isinstance(chat_history, list) else []
            
            try:
                # Get investigation result - FAST (uses fallback, no LLM)
                result = investigate_complete(message)
                
                if result.get("status") == "error":
                    ai_response = result.get("message", "Investigation failed")
                else:
                    ai_response = result.get("answer", "Investigation complete")
                    analysis = result.get("analysis", {})
                    
                    # Add quick summary
                    ai_response += f"\n\n**Summary:** {result.get('total_records', 0)} records | {analysis.get('unique_vehicles', 0)} vehicles | {analysis.get('unique_locations', 0)} locations"
                
                # SIMPLE: Just append assistant message dict
                assistant_msg_dict = {"role": "assistant", "content": ai_response}
                history.append(assistant_msg_dict)
                
                return history, result
            except Exception as e:
                print(f"Chat update error: {e}")
                import traceback
                traceback.print_exc()
                # Append error message
                history.append({"role": "assistant", "content": f"⚠️ Error: {str(e)[:100]}"})
                return history, (inv_results or {})
        
        def update_detailed_tabs_lightweight(inv_results):
            """Update detailed analysis tabs"""
            if not inv_results or inv_results.get("status") == "error":
                return ("No data available", "No findings yet", None, "No metrics")
            
            try:
                analysis = inv_results.get("analysis", {})
                total = inv_results.get('total_records', 0)
                
                # Summary
                data_info = f"📊 **{total} records** | 🚗 **{analysis.get('unique_vehicles', 0)} vehicles** | 📍 **{analysis.get('unique_locations', 0)} locations**"
                
                # Findings
                findings_list = analysis.get("key_findings", [])
                findings = "**Key Findings:**\n" + "\n".join([f"• {f}" for f in findings_list[:3]]) if findings_list else "No key findings"
                
                # Data table (limit to 20 rows)
                df_data = None
                try:
                    preview = inv_results.get("data_preview", [])[:20]
                    if preview:
                        df_data = pd.DataFrame(preview)
                except Exception as e:
                    print(f"DataFrame error: {e}")
                
                # Metrics
                metrics = f"**Confidence:** {analysis.get('confidence_score', 0):.0%}\n**Records:** {total}"
                
                return (data_info, findings, df_data, metrics)
            except Exception as e:
                print(f"Tab update error: {e}")
                return ("Error processing data", "Error", None, "Error")
        
        # OPTIMIZED CLICK HANDLER - Fast first update, then background refresh
        investigate_btn.click(
            fn=investigate_fast,
            inputs=[question_input, chatbot, conversation_state, investigation_results],
            outputs=[chatbot, conversation_state, investigation_results]
        ).then(
            fn=update_chat_final,
            inputs=[question_input, chatbot, investigation_results],
            outputs=[chatbot, investigation_results]
        ).then(
            fn=update_detailed_tabs_lightweight,
            inputs=[investigation_results],
            outputs=[data_structure, findings_display, data_table, metrics_display]
        ).then(
            fn=lambda inv: f"**📊 Stats**\n\n**Records:** {inv.get('total_records', 0)}\n**Vehicles:** {inv.get('analysis', {}).get('unique_vehicles', 0)}\n**Locations:** {inv.get('analysis', {}).get('unique_locations', 0)}" if inv.get("status") != "error" else "**Stats**\n\nAnalysis pending...",
            inputs=[investigation_results],
            outputs=[stats_display]
        ).then(
            fn=lambda inv: "\n".join([f"→ {q}" for q in inv.get("follow_ups", [])[:3]]) if inv.get("follow_ups") and inv.get("status") != "error" else "",
            inputs=[investigation_results],
            outputs=[suggested_follow_ups]
        ).then(
            fn=lambda: "",
            outputs=[question_input]
        )

    # =====================================================
    # TAB 4: Analytics
    # =====================================================

    with gr.Tab("Analytics", id="tab_analytics"):
        gr.Markdown("### 📊 Analytics Dashboard - Visual Intelligence")
        gr.Markdown("*Advanced analytics with real-time visualizations and data breakdown*")

        with gr.Row():
            refresh_btn = gr.Button(
                "🔄 Refresh Analytics",
                variant="primary",
                size="lg",
            )

        # ===== CHARTS ROW 1 =====
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                gr.Markdown("**🚗 Vehicles by State (Pie Chart)**")
                state_chart = gr.Image(
                    label="State Distribution",
                    type="pil"
                )

            with gr.Column(scale=1):
                gr.Markdown("**📊 Hourly Traffic Trends (Line Chart)**")
                hourly_chart = gr.Image(
                    label="Hourly Traffic",
                    type="pil"
                )

        # ===== CHARTS ROW 2 =====
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                gr.Markdown("**🏆 Top Detected Plates (Bar Chart)**")
                top_chart = gr.Image(
                    label="Top Plates",
                    type="pil"
                )

            with gr.Column(scale=1):
                gr.Markdown("**⚠️ Suspicious Vehicles (Donut Chart)**")
                suspicious_chart = gr.Image(
                    label="Suspicious Alerts",
                    type="pil"
                )

        # ===== DATA TABLES =====
        gr.Markdown("### 📋 Detailed Data Tables")

        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                state_table = gr.Dataframe(
                    label="Vehicles by State",
                    interactive=False,
                )

            with gr.Column(scale=1):
                hourly_table = gr.Dataframe(
                    label="Traffic by Hour",
                    interactive=False,
                )

        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                top_table = gr.Dataframe(
                    label="Top License Plates",
                    interactive=False,
                )

            with gr.Column(scale=1):
                suspicious_table = gr.Dataframe(
                    label="Suspicious Vehicles",
                    interactive=False,
                )

        refresh_btn.click(
            fn=refresh_analytics,
            outputs=[
                state_chart,
                hourly_chart,
                top_chart,
                suspicious_chart,
                state_table,
                hourly_table,
                top_table,
                suspicious_table
            ]
        )

    # =====================================================
    # VEHICLE INTELLIGENCE
    # =====================================================

    with gr.Tab("🗺️ Vehicle Intelligence"):

        vehicle_map = gr.HTML(
            value=default_vehicle_map()
        )

        with gr.Row():

            vehicle_plate = gr.Textbox(
                placeholder="TN57GR8753",
                label="License Plate (Required)",
                scale=2,
                info="Enter vehicle number plate"
            )

            date_from = gr.Textbox(
                label="Start Date (Optional)",
                placeholder="YYYY-MM-DD",
                scale=2,
                info="Format: YYYY-MM-DD or leave empty"
            )

            date_to = gr.Textbox(
                label="End Date (Optional)",
                placeholder="YYYY-MM-DD",
                scale=2,
                info="Format: YYYY-MM-DD or leave empty"
            )

            search_vehicle_btn = gr.Button(
                "🔍 Search Route",
                variant="primary",
                scale=1
            )

        with gr.Row():

            vehicle_table = gr.Dataframe(
                label="Detection Timeline",
                interactive=False
            )

            vehicle_info = gr.JSON(
                label="Route Analytics"
            )

        search_vehicle_btn.click(
            fn=search_vehicle_route,
            inputs=[
                vehicle_plate,
                date_from,
                date_to
            ],
            outputs=[
                vehicle_map,
                vehicle_table,
                vehicle_info
            ]
        )
# =========================================================
# QUEUE & LAUNCH
# =========================================================

demo.queue(max_size=20)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
        show_error=True,
    )