import streamlit as st
import numpy as np
import matplotlib.pyplot as plt

# App title and layout
st.set_page_config(page_title="AI Model Dashboard", layout="wide")

# Define models
models = [
    "Anomaly Detection",
    "Predictive Maintenance",
    "Cost Optimization",
    "Energy Efficiency"
]

# Initialize session state for agent status & errors
if "statuses" not in st.session_state:
    st.session_state.statuses = {model: "Running" for model in models}
if "error_logs" not in st.session_state:
    st.session_state.error_logs = {model: "" for model in models}
if "show_graphs" not in st.session_state:
    st.session_state.show_graphs = {model: False for model in models}

st.title("📊 AI Model & Agent Monitoring Dashboard")

# Function to handle status changes
def update_status(model, new_status):
    st.session_state.statuses[model] = new_status
    if new_status == "Error":
        st.session_state.error_logs[model] = f"Simulated error in {model} agent."
    elif new_status == "Resolved":
        st.session_state.show_graphs[model] = True
    elif new_status == "Running":
        st.session_state.show_graphs[model] = False

# Layout for models
for model in models:
    col1, col2, col3 = st.columns([2, 1, 1])
    
    # Model & Agent Status
    with col1:
        st.subheader(f"⚙️ {model} - Agent Status")
    
    # Display Status
    with col2:
        status = st.session_state.statuses[model]
        if status == "Running":
            st.success("✅ Running - No Issues")
        elif status == "Error":
            st.error("❌ Error Occurred")
        elif status == "Resolving":
            st.warning("🔄 Resolving Issue...")
        elif status == "Resolved":
            st.success("✔ Resolved & Completed")

    # Control Buttons
    with col3:
        if status == "Running":
            if st.button(f"🛑 Simulate Error {model}", key=f"error_{model}"):
                update_status(model, "Error")
        elif status == "Error":
            if st.button(f"🔄 Resolve & Restart {model}", key=f"resolve_{model}"):
                update_status(model, "Resolving")
        elif status == "Resolving":
            if st.button(f"✔ Mark Resolved {model}", key=f"mark_resolved_{model}"):
                update_status(model, "Resolved")
        elif status == "Resolved":
            if st.button(f"▶ Restart {model}", key=f"restart_{model}"):
                update_status(model, "Running")

    # Show error logs if any
    if st.session_state.error_logs[model]:
        st.error(f"Error Log: {st.session_state.error_logs[model]}")

    # Display graph when resolved
    if st.session_state.show_graphs[model]:
        st.subheader(f"📊 {model} - Performance Graph")
        fig, ax = plt.subplots()
        x = np.arange(1, 11)
        y = np.random.rand(10) * 100  # Random performance data
        ax.plot(x, y, marker="o", linestyle="-", color="b")
        ax.set_xlabel("Time")
        ax.set_ylabel("Performance")
        ax.set_title(f"{model} - Performance Over Time")
        st.pyplot(fig)

    st.divider()

st.info("⚠ If an error occurs, only the affected agent stops. Resolve it before moving to the next.")