import hashlib
from datetime import datetime

import gradio as gr
import pandas as pd
import plotly.graph_objects as go
from datasets import load_dataset

DATASET_REPO_PATH = "uclanecl/NECL_GPUs"
USER_BLACKLIST = ["gdm", "necl", "ncel"]


def get_plot_df(df):
    df = df.fillna({"fan_speed": "N/A"})

    df["users"] = df["users"].str.split(",")

    # Filter out blacklisted users
    df["users"] = df["users"].apply(
        lambda user_list: [u for u in user_list if u not in USER_BLACKLIST] or ["Free"]
    )

    # Construct a new DataFrame for the bar chart
    new_df = []
    for _, row_series in df.iterrows():
        users_list = row_series["users"]
        gpu_users_num = len(row_series["users"])

        for username in users_list:
            new_row = row_series.copy()
            new_row = new_row.drop(labels="users")

            new_row["count"] = 1 / gpu_users_num
            new_row["username"] = username

            new_df.append(new_row)

    df = pd.DataFrame(new_df)

    return df


def username_to_hsl(username):
    # Black indicates a free GPU
    if username == "Free":
        return "hsl(0, 0%, 0%)"

    # Deterministic hash → integer
    n = int.from_bytes(hashlib.sha1(username.encode()).digest(), "big")

    # Map to hue [0, 360). The 15 offset is to customize the color for a "specific user.""
    hue = (n + 15) % 360

    # Fixed saturation and lightness
    saturation = 50  # percent
    lightness = 80  # percent

    # Return as HSL CSS-style string
    return f"hsl({hue}, {saturation}%, {lightness}%)"


def get_bar_chart(df):
    # Create the bar chart
    fig = go.Figure()

    for _, row_series in df.iterrows():
        # Using the low-level interface because the high-level interface only allows contiguous
        # colored blocks in one bar
        fig.add_bar(
            x=[row_series["count"]],
            y=[row_series["server"]],
            orientation="h",
            marker_color=username_to_hsl(row_series["username"]),
            text=f"{row_series['index']}<br>{row_series['username']}",
            textposition="inside",
            insidetextanchor="middle",
            hoverinfo="none",
        )

    last_update_time = max(df["query_time"])
    last_update_time = datetime.fromisoformat(last_update_time).strftime("%Y-%m-%d %H:%M:%S")

    max_index = df["index"].max()

    fig.update_layout(
        barmode="stack",
        title=f"Last Modified {last_update_time}",
        showlegend=False,
        modebar_remove=["lasso2d", "select2d"],
        height=500,
        plot_bgcolor="#E4E4E6",
        xaxis_fixedrange=True,
        xaxis_range=[0, max_index + 1],
        xaxis_showticklabels=False,
        yaxis_fixedrange=True,
        yaxis_categoryorder="category descending",
    )

    return fig


def get_interface_elements():
    try:
        df = load_dataset(DATASET_REPO_PATH)["train"].to_pandas()
        df = get_plot_df(df)

        plot = get_bar_chart(df)

        # Reorganize DataFrame that will be displayed
        display_df_columns = [
            "server",
            "index",
            "username",
            "device",
            "memory",
            "memory_usage",
            "utilization",
            "temperature",
            "fan_speed",
            "query_time",
        ]
        display_df = df[display_df_columns]

        return plot, display_df

    except Exception as e:
        print(f"Global error in plot: {e}")
        return None, pd.DataFrame()


demo = gr.Interface(
    fn=get_interface_elements,
    inputs=[],
    outputs=[
        gr.Plot(label="GPU Status", elem_classes="gpu_status_plot"),
        gr.Dataframe(label="GPU Status Details"),
    ],
    live=True,
    css=".gpu_status_plot {max-width: 800px;}",
)

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