"""Admin Dashboard - Real-time monitoring and control interface."""
import streamlit as st
import time
from datetime import datetime
from typing import Dict, List, Any
import asyncio


# Page configuration
st.set_page_config(
    page_title="LogiFlow AI - Agent Control Center",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .agent-card {
        background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
        border-radius: 12px;
        padding: 20px;
        margin: 10px 0;
        border-left: 4px solid #00d4ff;
    }
    .status-badge {
        display: inline-block;
        padding: 4px 12px;
        border-radius: 20px;
        font-size: 12px;
        font-weight: bold;
    }
    .status-idle { background: #4ade80; color: #000; }
    .status-working { background: #fbbf24; color: #000; }
    .status-error { background: #f87171; color: #fff; }
    .metric-card {
        background: #1e1e1e;
        border-radius: 8px;
        padding: 15px;
        text-align: center;
    }
    .metric-value { font-size: 32px; font-weight: bold; color: #00d4ff; }
    .metric-label { font-size: 14px; color: #888; }
    .order-card {
        background: #252525;
        border-radius: 8px;
        padding: 15px;
        margin: 8px 0;
    }
</style>
""", unsafe_allow_html=True)


class DashboardState:
    """Shared state for dashboard."""
    
    def __init__(self):
        self.agents = {
            "JARVIS": {"status": "idle", "tasks_completed": 0, "current_task": None},
            "ORDER_AGENT": {"status": "idle", "tasks_completed": 0, "current_task": None},
            "DISPATCH_AGENT": {"status": "idle", "tasks_completed": 0, "current_task": None},
            "FLEET_AGENT": {"status": "idle", "tasks_completed": 0, "current_task": None},
            "HUB_AGENT": {"status": "idle", "tasks_completed": 0, "current_task": None},
            "CUSTOMER_AGENT": {"status": "idle", "tasks_completed": 0, "current_task": None},
            "EXCEPTION_AGENT": {"status": "idle", "tasks_completed": 0, "current_task": None},
            "ROUTE_OPTIMIZER": {"status": "idle", "tasks_completed": 0, "current_task": None},
        }
        self.orders = []
        self.vehicles = []
        self.hubs = []
        self.messages = []
        
    def update_agent_status(self, agent_id: str, status: str, task: str = None):
        if agent_id in self.agents:
            self.agents[agent_id]["status"] = status
            self.agents[agent_id]["current_task"] = task
            if status == "idle" and task is None:
                self.agents[agent_id]["tasks_completed"] += 1
    
    def add_message(self, sender: str, recipient: str, message_type: str, content: str):
        self.messages.append({
            "timestamp": datetime.now(),
            "sender": sender,
            "recipient": recipient,
            "type": message_type,
            "content": content
        })
        if len(self.messages) > 100:
            self.messages = self.messages[-100:]


# Initialize state
if 'state' not in st.session_state:
    st.session_state.state = DashboardState()


def render_header():
    """Render the dashboard header."""
    col1, col2, col3 = st.columns([2, 3, 1])
    
    with col1:
        st.markdown("""
        <div style="padding: 20px;">
            <h1 style="margin: 0; color: #00d4ff;">🤖 LogiFlow AI</h1>
            <p style="margin: 5px 0 0 0; color: #888;">Agent Control Center</p>
        </div>
        """, unsafe_allow_html=True)
    
    with col2:
        st.markdown("""
        <div style="padding: 20px; text-align: center;">
            <h3 style="color: #4ade80;">🟢 System Operational</h3>
            <p style="color: #888; margin: 0;">All agents active and monitoring</p>
        </div>
        """, unsafe_allow_html=True)
    
    with col3:
        st.markdown(f"""
        <div style="padding: 20px; text-align: right;">
            <p style="color: #888; margin: 0;">{datetime.now().strftime('%Y-%m-%d')}</p>
            <p style="color: #00d4ff; font-size: 24px; margin: 5px 0;">{datetime.now().strftime('%H:%M:%S')}</p>
        </div>
        """, unsafe_allow_html=True)


def render_agent_grid(state: DashboardState):
    """Render agent status cards."""
    st.markdown("## 🤖 Agent Status")
    
    cols = st.columns(4)
    agent_icons = {
        "JARVIS": "🧠",
        "ORDER_AGENT": "📦",
        "DISPATCH_AGENT": "🚚",
        "FLEET_AGENT": "🚛",
        "HUB_AGENT": "🏭",
        "CUSTOMER_AGENT": "💬",
        "EXCEPTION_AGENT": "⚠️",
        "ROUTE_OPTIMIZER": "🗺️"
    }
    
    for i, (agent_id, info) in enumerate(state.agents.items()):
        with cols[i % 4]:
            status_class = f"status-{info['status']}"
            icon = agent_icons.get(agent_id, "🤖")
            
            st.markdown(f"""
            <div class="agent-card">
                <div style="display: flex; justify-content: space-between; align-items: center;">
                    <h4 style="margin: 0;">{icon} {agent_id}</h4>
                    <span class="status-badge {status_class}">{info['status'].upper()}</span>
                </div>
                <div style="margin-top: 10px; color: #888;">
                    <p style="margin: 5px 0;">Tasks: {info['tasks_completed']}</p>
                    <p style="margin: 5px 0;">Current: {info['current_task'] or 'None'}</p>
                </div>
            </div>
            """, unsafe_allow_html=True)


def render_metrics(state: DashboardState):
    """Render key metrics."""
    col1, col2, col3, col4, col5 = st.columns(5)
    
    metrics = [
        ("Active Orders", len([o for o in state.orders if o.get("status") != "delivered"]), "📦"),
        ("Vehicles Active", len([v for v in state.vehicles if v.get("status") == "in_transit"]), "🚛"),
        ("Hubs Online", len(state.hubs), "🏭"),
        ("Messages/min", len([m for m in state.messages if (datetime.now() - m["timestamp"]).seconds < 60]), "📨"),
        ("Success Rate", "98.5%", "✅")
    ]
    
    for col, (label, value, icon) in zip([col1, col2, col3, col4, col5], metrics):
        with col:
            st.markdown(f"""
            <div class="metric-card">
                <div style="font-size: 32px;">{icon}</div>
                <div class="metric-value">{value}</div>
                <div class="metric-label">{label}</div>
            </div>
            """, unsafe_allow_html=True)


def render_order_list(state: DashboardState):
    """Render recent orders."""
    st.markdown("## 📦 Recent Orders")
    
    if not state.orders:
        st.info("No orders yet. Create an order to see it here.")
        return
    
    for order in state.orders[-10:][::-1]:
        status_colors = {
            "received": "🟡",
            "in_transit": "🔵",
            "delivered": "🟢",
            "cancelled": "🔴"
        }
        status_icon = status_colors.get(order.get("status", "received"), "⚪")
        
        with st.expander(f"{status_icon} {order.get('order_id', 'N/A')} - {order.get('customer_name', 'Unknown')}"):
            col1, col2 = st.columns(2)
            with col1:
                st.write(f"**Status:** {order.get('status', 'Unknown')}")
                st.write(f"**Priority:** {order.get('category', {}).get('priority', 'Medium')}")
                st.write(f"**Created:** {order.get('created_at', 'N/A')}")
            with col2:
                st.write(f"**Pickup:** {order.get('pickup_address', {}).get('city', 'N/A')}")
                st.write(f"**Delivery:** {order.get('delivery_address', {}).get('city', 'N/A')}")
                st.write(f"**Items:** {len(order.get('items', []))}")


def render_vehicle_fleet(state: DashboardState):
    """Render vehicle fleet status."""
    st.markdown("## 🚛 Vehicle Fleet")
    
    if not state.vehicles:
        st.info("No vehicles registered.")
        return
    
    for vehicle in state.vehicles[:20]:
        status_icons = {
            "available": "🟢",
            "in_transit": "🔵",
            "maintenance": "🟡",
            "offline": "⚫"
        }
        icon = status_icons.get(vehicle.get("status", "available"), "⚪")
        
        col1, col2, col3, col4 = st.columns(4)
        with col1:
            st.write(f"{icon} **{vehicle.get('vehicle_id', 'N/A')}**")
        with col2:
            st.write(f"Type: {vehicle.get('type', 'N/A')}")
        with col3:
            st.write(f"Hub: {vehicle.get('hub', 'N/A')}")
        with col4:
            st.write(f"Status: {vehicle.get('status', 'N/A')}")
        st.divider()


def render_hub_network(state: DashboardState):
    """Render hub network status."""
    st.markdown("## 🏭 Hub Network")
    
    cols = st.columns(4)
    hub_data = [
        {"id": "DL-HUB-01", "name": "Delhi North", "load": 150, "capacity": 500},
        {"id": "DL-HUB-02", "name": "Delhi South", "load": 200, "capacity": 400},
        {"id": "MU-HUB-01", "name": "Mumbai West", "load": 300, "capacity": 600},
        {"id": "MU-HUB-02", "name": "Mumbai East", "load": 250, "capacity": 550},
        {"id": "BL-HUB-01", "name": "Bangalore", "load": 180, "capacity": 500},
        {"id": "HY-HUB-01", "name": "Hyderabad", "load": 220, "capacity": 450},
        {"id": "PU-HUB-01", "name": "Pune", "load": 160, "capacity": 400},
        {"id": "KL-HUB-01", "name": "Kolkata", "load": 140, "capacity": 350},
    ]
    
    for i, hub in enumerate(hub_data):
        with cols[i % 4]:
            utilization = (hub["load"] / hub["capacity"]) * 100
            color = "🟢" if utilization < 70 else "🟡" if utilization < 90 else "🔴"
            
            st.markdown(f"""
            <div class="agent-card">
                <h4>{color} {hub['name']}</h4>
                <p style="margin: 5px 0;">ID: {hub['id']}</p>
                <p style="margin: 5px 0;">Load: {hub['load']}/{hub['capacity']}</p>
                <p style="margin: 5px 0; color: #00d4ff;">Utilization: {utilization:.1f}%</p>
            </div>
            """, unsafe_allow_html=True)


def render_message_feed(state: DashboardState):
    """Render real-time message feed."""
    st.markdown("## 📨 Agent Communication Feed")
    
    if not state.messages:
        st.info("No messages yet.")
        return
    
    for msg in state.messages[-20:][::-1]:
        st.markdown(f"""
        <div style="background: #1e1e1e; padding: 10px; border-radius: 5px; margin: 5px 0;">
            <span style="color: #888;">{msg['timestamp'].strftime('%H:%M:%S')}</span>
            <span style="color: #00d4ff;"> [{msg['sender']}]</span>
            → <span style="color: #4ade80;">[{msg['recipient']}]</span>
            <span style="color: #fbbf24;"> {msg['type']}</span>
            <p style="margin: 5px 0 0 0; color: #ccc;">{msg['content']}</p>
        </div>
        """, unsafe_allow_html=True)


def render_order_creator():
    """Render order creation form."""
    st.markdown("## ➕ Create New Order")
    
    with st.form("order_form"):
        col1, col2 = st.columns(2)
        
        with col1:
            st.write("**Customer Details**")
            customer_name = st.text_input("Customer Name", "Rajesh Kumar")
            customer_phone = st.text_input("Phone", "+919876543210")
            customer_email = st.text_input("Email", "rajesh@example.com")
        
        with col2:
            st.write("**Order Priority**")
            priority = st.selectbox("Priority", ["low", "medium", "high", "urgent"])
            items_count = st.number_input("Number of Items", 1, 50, 3)
            weight = st.number_input("Total Weight (kg)", 0.1, 500.0, 5.0)
        
        st.write("**Pickup Address**")
        col_p1, col_p2, col_p3 = st.columns(3)
        with col_p1:
            pickup_city = st.text_input("City", "Delhi")
        with col_p2:
            pickup_state = st.text_input("State", "Delhi")
        with col_p3:
            pickup_pincode = st.text_input("Pincode", "110001")
        
        st.write("**Delivery Address**")
        col_d1, col_d2, col_d3 = st.columns(3)
        with col_d1:
            delivery_city = st.text_input("City", "Mumbai", key="del_city")
        with col_d2:
            delivery_state = st.text_input("State", "Maharashtra", key="del_state")
        with col_d3:
            delivery_pincode = st.text_input("Pincode", "400001", key="del_pin")
        
        submitted = st.form_submit_button("🚀 Create Order via AI Agents", use_container_width=True)
        
        if submitted:
            st.success("Order created! Watch the agents process it in real-time.")
            
            # Simulate order being created
            if 'state' in st.session_state:
                state = st.session_state.state
                order_id = f"ORD-{datetime.now().strftime('%Y%m%d')}-{len(state.orders) + 1:04d}"
                
                state.orders.append({
                    "order_id": order_id,
                    "customer_name": customer_name,
                    "customer_phone": customer_phone,
                    "status": "received",
                    "priority": priority,
                    "pickup_address": {"city": pickup_city, "pincode": pickup_pincode},
                    "delivery_address": {"city": delivery_city, "pincode": delivery_pincode},
                    "items": [{"name": f"Item {i}", "weight": weight/items_count} for i in range(int(items_count))],
                    "created_at": datetime.now().isoformat(),
                    "category": {"priority": priority}
                })
                
                state.add_message("ADMIN", "JARVIS", "ORDER_SUBMIT", f"New order {order_id} submitted")
                state.update_agent_status("ORDER_AGENT", "working", f"Validating {order_id}")
                
            st.rerun()


def render_exception_panel(state: DashboardState):
    """Render exception handling panel."""
    st.markdown("## ⚠️ Exception Management")
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.write("**Quick Actions**")
        if st.button("🚨 Simulate Delay Exception"):
            state.add_message("SYSTEM", "EXCEPTION_AGENT", "EXCEPTION", "Simulated delay detected")
            st.warning("Delay exception triggered - Exception Agent handling...")
    
    with col2:
        st.write("**Exception Queue**")
        st.info("No active exceptions - System running smoothly")


def main():
    """Main dashboard function."""
    
    # Sidebar navigation
    st.sidebar.markdown("## 🧭 Navigation")
    page = st.sidebar.radio("Go to", [
        "📊 Overview",
        "🤖 Agents",
        "📦 Orders",
        "🚛 Fleet",
        "🏭 Hubs",
        "⚠️ Exceptions",
        "💬 Messages",
        "➕ Create Order"
    ])
    
    # Initialize state
    state = st.session_state.state
    
    # Simulate some data if empty
    if not state.vehicles:
        state.vehicles = [
            {"vehicle_id": "DL-V-001", "type": "van", "hub": "DL-HUB-01", "status": "in_transit", "location": {"lat": 28.61, "lng": 77.20}},
            {"vehicle_id": "DL-V-002", "type": "van", "hub": "DL-HUB-01", "status": "available", "location": {"lat": 28.62, "lng": 77.21}},
            {"vehicle_id": "MU-V-001", "type": "van", "hub": "MU-HUB-01", "status": "in_transit", "location": {"lat": 19.07, "lng": 72.87}},
            {"vehicle_id": "BL-V-001", "type": "van", "hub": "BL-HUB-01", "status": "available", "location": {"lat": 12.97, "lng": 77.59}},
            {"vehicle_id": "DL-B-001", "type": "bike", "hub": "DL-HUB-01", "status": "in_transit", "location": {"lat": 28.63, "lng": 77.19}},
        ]
    
    if not state.hubs:
        state.hubs = [
            {"hub_id": "DL-HUB-01", "name": "Delhi North Hub"},
            {"hub_id": "DL-HUB-02", "name": "Delhi South Hub"},
            {"hub_id": "MU-HUB-01", "name": "Mumbai West Hub"},
            {"hub_id": "BL-HUB-01", "name": "Bangalore Hub"},
        ]
    
    # Render based on selection
    if page == "📊 Overview":
        render_header()
        render_metrics(state)
        render_agent_grid(state)
        
        col_overview1, col_overview2 = st.columns(2)
        with col_overview1:
            render_order_list(state)
        with col_overview2:
            render_hub_network(state)
            
    elif page == "🤖 Agents":
        render_header()
        st.markdown("# 🤖 Agent Control Center")
        render_agent_grid(state)
        
        st.markdown("### Agent Communication Pattern")
        st.markdown("""
        ```
        JARVIS (Master)
            ├── ORDER_AGENT → Validates & categorizes orders
            ├── DISPATCH_AGENT → Assigns routes & zones
            ├── FLEET_AGENT → Manages vehicle allocation
            ├── HUB_AGENT → Coordinates hub operations
            ├── CUSTOMER_AGENT → Sends notifications
            ├── EXCEPTION_AGENT → Handles problems
            └── ROUTE_OPTIMIZER → Plans optimal paths
        ```
        """)
        
    elif page == "📦 Orders":
        render_header()
        render_order_list(state)
        render_order_creator()
        
    elif page == "🚛 Fleet":
        render_header()
        render_vehicle_fleet(state)
        
    elif page == "🏭 Hubs":
        render_header()
        render_hub_network(state)
        
    elif page == "⚠️ Exceptions":
        render_header()
        render_exception_panel(state)
        
    elif page == "💬 Messages":
        render_header()
        render_message_feed(state)
        
    elif page == "➕ Create Order":
        render_header()
        render_order_creator()
    
    # Auto-refresh
    time.sleep(1)
    st.rerun()


if __name__ == "__main__":
    main()