ui.py

import streamlit as st
import requests
import pandas as pd
from datetime import datetime
from streamlit_autorefresh import st_autorefresh

API_URL = "http://localhost:5000"

# State for tracking CPU over time
if "cpu_chart_data" not in st.session_state:
    st.session_state.cpu_chart_data = pd.DataFrame(columns=["Time", "Node", "Available CPU"])

def status_icon(status):
    return "🟢" if status == "UP" else "🔴"

def fetch_nodes():
    try:
        res = requests.get(f"{API_URL}/list_nodes")
        res.raise_for_status()
        nodes = res.json()
        if not nodes:
            return pd.DataFrame(), pd.DataFrame()

        # Main table
        df = pd.DataFrame(nodes)
        df['status_icon'] = df['status'].apply(status_icon)
        df['last_heartbeat'] = pd.to_datetime(df['last_heartbeat'], unit='s').dt.strftime('%Y-%m-%d %H:%M:%S')

        table = df[["node_id", "cpu_total", "cpu_available", "status_icon", "last_heartbeat", "pods"]].rename(columns={
            'node_id': 'ID',
            'cpu_total': 'CPU Cores',
            'cpu_available': 'Available CPU',
            'status_icon': 'Status',
            'last_heartbeat': 'Last Heartbeat',
            'pods': 'Pods'
        })

        # Update chart history
        now = datetime.now().strftime("%H:%M:%S")
        chart_df = pd.DataFrame([{
            "Time": now,
            "Node": row["node_id"],
            "Available CPU": row["cpu_available"]
        } for _, row in df.iterrows()])
        st.session_state.cpu_chart_data = pd.concat([st.session_state.cpu_chart_data, chart_df], ignore_index=True)

        return table, st.session_state.cpu_chart_data

    except Exception as e:
        st.error(f"Failed to fetch nodes: {str(e)}")
        return pd.DataFrame(), pd.DataFrame()

def fetch_pods():
    try:
        res = requests.get(f"{API_URL}/list_pods")
        res.raise_for_status()
        pods = res.json()
        if not pods:
            return pd.DataFrame()

        df = pd.DataFrame(pods)
        df['created_at'] = pd.to_datetime(df['created_at']).dt.strftime('%Y-%m-%d %H:%M:%S')
        return df[["pod_id", "node_id", "cpu", "status", "created_at"]].rename(columns={
            'pod_id': 'ID',
            'node_id': 'Node',
            'cpu': 'CPU',
            'status': 'Status',
            'created_at': 'Created At'
        })

    except Exception as e:
        st.error(f"Failed to fetch pods: {str(e)}")
        return pd.DataFrame()

def add_node():
    with st.form("add_node_form"):
        st.write("### Add New Node")
        node_id = st.text_input("Node ID")
        cpu_cores = st.number_input("CPU Cores", min_value=1, value=4)
        if st.form_submit_button("Add Node") and node_id:
            response = requests.post(
                f"{API_URL}/add_node",
                json={"node_id": node_id, "cpu_cores": cpu_cores}
            )
            if response.status_code == 201:
                st.success(response.json().get("message", "Node added successfully"))
            else:
                st.error(response.json().get("error", "Failed to add node"))

def launch_pod():
    with st.form("launch_pod_form"):
        st.write("### Launch New Pod")
        pod_id = st.text_input("Pod ID")
        cpu_required = st.number_input("CPU Required", min_value=1, value=1)
        if st.form_submit_button("Launch Pod") and pod_id:
            response = requests.post(
                f"{API_URL}/launch_pod",
                json={"pod_id": pod_id, "cpu": cpu_required}
            )
            if response.status_code == 201:
                st.success(response.json().get("message", "Pod launched successfully"))
            else:
                st.error(response.json().get("error", "Failed to launch pod"))

# UI Setup
st.set_page_config(page_title="Kawaii-Kube", layout="wide")
st.title("🌸 Kawaii-Kube: Cluster Simulator")

# Layout tabs
tab1, tab2 = st.tabs(["⚙️ Cluster Management", "📊 Monitoring"])

# Refresh every 10s
st_autorefresh(interval=10_000, key="refresh")

with tab1:
    col1, col2 = st.columns(2)
    with col1:
        add_node()
    with col2:
        launch_pod()

with tab2:
    st.write("### Cluster Nodes")
    table_df, chart_df = fetch_nodes()
    st.dataframe(table_df, use_container_width=True)

    st.write("### CPU Availability Over Time")
    if not chart_df.empty:
        pivot_chart = chart_df.pivot(index="Time", columns="Node", values="Available CPU").fillna(method="ffill")
        st.line_chart(pivot_chart)

    st.write("### Running Pods")
    st.dataframe(fetch_pods(), use_container_width=True)

    st.caption(f"Last updated: {datetime.now().strftime('%H:%M:%S')}")