import os
import gradio as gr
import numpy as np
import pandas as pd
import datetime
import plotly.express as px
import plotly.graph_objects as go
import datasets


##### GPU PLOT #####
def split_multi_users(dfs):
    df = dfs.copy()
    df["usernames"] = df["username"].apply(lambda x: x.split(", "))
    df["count"] = 1
    new_df = []
    for row in df.to_dict(orient="records"):
        gpu_users_num = len(row["usernames"])
        for username in row["usernames"]:
            new_row = row.copy()
            new_row["count"] = 1 / gpu_users_num
            new_row["username"] = username
            new_df.append(new_row)
    df = pd.DataFrame(new_df)
    return df

def plot_gpus():
    ### Load Data
    dfs = datasets.load_dataset("pluslab/PLUS_Lab_GPUs_Data", data_files="gpus.csv", download_mode='force_redownload')["train"].to_pandas()
    dfs = dfs.drop(columns=["Unnamed: 0"])
    dfs = dfs.fillna("FREE")
    dfs_plot = split_multi_users(dfs)
    fig = px.bar(
        dfs_plot, x="count", y="server", color="username",
        title=f"Last Updated {min(dfs['timestamp'])}",
        color_discrete_map={
            "FREE": "black",
        },
        text=dfs_plot['username'].astype(str) + "<br>" + dfs_plot['device'].astype(str),
    )
    fig.update_layout(
        yaxis={'categoryorder': 'array', 'categoryarray': dfs_plot["server"].unique()[::-1]},
        barcornerradius=5,
    )
    fig.update_traces(textposition='inside', insidetextanchor='middle')
    # print(dfs_plot)
    return fig, dfs

##### DISK PLOT #####
def _pick_col(df, candidates):
    norm = {c.strip().lower(): c for c in df.columns}
    for cand in candidates:
        cand = cand.strip().lower()
        if cand in norm:
            return norm[cand]
    return None

def _kblocks_to_tib(kblocks):  # shown as "TB" per your convention
    return kblocks / (1024**3)

def _kblocks_to_gib(kblocks):  # shown as "GB"
    return kblocks / (1024**2)

def plot_disks(alert_threshold_pct=99.0):
    df = datasets.load_dataset(
        "pluslab/PLUS_Lab_GPUs_Data",
        data_files="disks.csv",
        download_mode="force_redownload",
    )["train"].to_pandas()

    if "Unnamed: 0" in df.columns:
        df = df.drop(columns=["Unnamed: 0"])

    df = pd.concat([df[df['Mounted'] != '/data2'], df[df['Mounted'] == '/data2'].drop_duplicates(subset=['Mounted'])])  # Keep one of /data2

    server_col = _pick_col(df, ["server"])
    fs_col     = _pick_col(df, ["filesystem"])
    blocks_col = _pick_col(df, ["1k-blocks", "1k blocks", "blocks"])
    used_col   = _pick_col(df, ["used"])
    avail_col  = _pick_col(df, ["available", "avail"])
    mount_col  = _pick_col(df, ["mounted", "mounted on", "mount", "mountpoint"])

    required = [server_col, fs_col, blocks_col, used_col, avail_col]
    if any(c is None for c in required):
        raise ValueError(f"Missing required columns. Found: {list(df.columns)}")

    for c in [blocks_col, used_col, avail_col]:
        df[c] = pd.to_numeric(df[c], errors="coerce")

    # Label
    if mount_col is not None:
        df["Label"] = df[server_col].astype(str) + " • " + df[mount_col].astype(str)
    else:
        df["Label"] = df[server_col].astype(str) + " • " + df[fs_col].astype(str)

    # Totals & pct
    df["Total_kb"] = df[used_col] + df[avail_col]
    df["Used_pct"] = (df[used_col] / df["Total_kb"]) * 100.0
    df["Used_pct"] = df["Used_pct"].clip(0, 100)
    df["Avail_pct"] = (100.0 - df["Used_pct"]).clip(0, 100)

    # Sizes
    df["Used_TB"]  = _kblocks_to_tib(df[used_col])
    df["Total_TB"] = _kblocks_to_tib(df["Total_kb"])
    df["Avail_GB"] = _kblocks_to_gib(df[avail_col])  # <-- GB for hovers

    # Alerts
    df["ALERT"] = df["Used_pct"] > alert_threshold_pct

    # Sort
    # df = df.sort_values("Total_kb", ascending=False).reset_index(drop=True)
    df = df.sort_values("Mounted", ascending=False).reset_index(drop=True)

    y_labels = np.where(df["ALERT"].to_numpy(), "⚠ " + df["Label"], df["Label"])

    used_text   = [f"{u:.2f} TB ({p:.0f}%)" for u, p in zip(df["Used_TB"], df["Used_pct"])]
    total_annot = [f"{t:.2f} TB" for t in df["Total_TB"]]
    avail_gb_0  = [f"{g:.0f} GB" for g in df["Avail_GB"]]

    # Colors
    COLOR_TOTAL = "#CBD5E1"
    COLOR_USED  = "#2563EB"
    COLOR_FREE  = "#94A3B8"
    COLOR_ALERT = "#F59E0B"
    used_colors = np.where(df["ALERT"].to_numpy(), COLOR_ALERT, COLOR_USED)

    fig = go.Figure()

    # Gray background hover: Available in GB (0dp)
    fig.add_trace(
        go.Bar(
            y=y_labels,
            x=[100] * len(df),
            base=0,
            orientation="h",
            marker=dict(color=COLOR_TOTAL),
            opacity=0.40,
            hovertemplate="<b>%{y}</b><br>Available: %{customdata}<br><extra></extra>",
            customdata=avail_gb_0,
            showlegend=False,
        )
    )

    # Used hover: Available in GB (0dp) too
    fig.add_trace(
        go.Bar(
            y=y_labels,
            x=df["Used_pct"],
            base=0,
            name=f"Used (>{alert_threshold_pct:.0f}% highlighted)",
            orientation="h",
            marker=dict(color=used_colors),
            text=used_text,
            textposition="inside",
            insidetextanchor="middle",
            hovertemplate=(
                "<b>%{y}</b><br>"
                "Used: %{customdata[0]} (%{customdata[3]:.2f}%)<br>"
                "Available: %{customdata[1]}<br>"
                "Total: %{customdata[2]}<br>"
                "<extra></extra>"
            ),
            customdata=np.stack(
                [
                    df["Used_TB"].map(lambda v: f"{v:.2f} TB").to_numpy(),
                    df["Avail_GB"].map(lambda v: f"{v:.0f} GB").to_numpy(),  # <-- changed
                    df["Total_TB"].map(lambda v: f"{v:.2f} TB").to_numpy(),
                    df["Used_pct"].to_numpy(),
                ],
                axis=1,
            ),
        )
    )

    # Available hover: Available in GB (0dp)
    fig.add_trace(
        go.Bar(
            y=y_labels,
            x=df["Avail_pct"],
            base=df["Used_pct"],
            name="Available",
            orientation="h",
            marker=dict(color=COLOR_FREE),
            hovertemplate=(
                "<b>%{y}</b><br>"
                "Available: %{customdata[0]}<br>"
                "Used: %{customdata[1]}<br>"
                "Total: %{customdata[2]}<br>"
                "<extra></extra>"
            ),
            customdata=np.stack(
                [
                    df["Avail_GB"].map(lambda v: f"{v:.0f} GB").to_numpy(),  # <-- changed
                    df["Used_TB"].map(lambda v: f"{v:.2f} TB").to_numpy(),
                    df["Total_TB"].map(lambda v: f"{v:.2f} TB").to_numpy(),
                ],
                axis=1,
            ),
        )
    )

    # Total annotation (TB, 2dp)
    for y, ttxt, is_alert in zip(y_labels, total_annot, df["ALERT"].to_numpy()):
        fig.add_annotation(
            x=100,
            y=y,
            text=ttxt,
            showarrow=False,
            xanchor="left",
            yanchor="middle",
            xshift=6,
            font=dict(color=("#B45309" if is_alert else "#334155")),
        )

    fig.update_layout(
        barmode="overlay",
        template="plotly_white",
        title=f"Disk usage (alerts: Used > {alert_threshold_pct:.0f}%)",
        xaxis=dict(range=[0, 100], ticksuffix="%", title="Percent of total"),
        yaxis_title="",
        height=max(420, 28 * len(df)),
        margin=dict(l=280, r=120, t=60, b=40),
        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        barcornerradius=4,
    )
    fig.update_yaxes(autorange="reversed")
    return fig, df



##### PLOT ALL #####
def plot_figs():
    fig_gpus, df_gpus = plot_gpus()
    fig_disks, df_disks = plot_disks()
    return fig_gpus, fig_disks, df_gpus, df_disks

demo = gr.Interface(
    fn=plot_figs,
    inputs = [

    ],
    outputs = [
        gr.Plot(label="GPU Status", elem_classes="plotcss"),
        gr.Plot(label="Disk Status", elem_classes="plotcss"),
        gr.Dataframe(label="GPU Status Details"),
        gr.Dataframe(label="Disk Status Details"),
    ],
    live=True,
    flagging_options=[],
    css=".plotcss {max-width: 820px !important;}",
    delete_cache=(1, 1)
)

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