import json
import os
import time
from dataclasses import dataclass
from datetime import datetime
from zoneinfo import ZoneInfo

import gradio as gr
import plotly.graph_objects as go
import wandb
from substrateinterface import Keypair
from wandb.apis.public import Run
import numpy as np

WANDB_RUN_PATH = os.environ["WANDB_RUN_PATH"]
SOURCE_VALIDATOR_UID = int(os.environ["SOURCE_VALIDATOR_UID"])
BASELINE_AVERAGE = float(os.environ["BASELINE_AVERAGE"])

REFRESH_RATE = 60 * 30  # 30 minutes
GRAPH_HISTORY_DAYS = 30
MAX_GRAPH_ENTRIES = 10

wandb_api = wandb.Api()
demo = gr.Blocks(css=".typewriter {font-family: 'JMH Typewriter', sans-serif;}")
runs: dict[int, list[Run]] = {}


@dataclass
class LeaderboardEntry:
    uid: int
    model: str
    score: float
    model_average: float
    similarity: float
    hotkey: str
    previous_day_winner: bool
    rank: int


@dataclass
class GraphEntry:
    dates: list[datetime]
    generation_times: list[float]
    similarities: list[float]
    scores: list[float]
    models: list[str]
    best_time: float


def is_valid_run(run: Run):
    required_config_keys = ["hotkey", "uid", "contest", "signature"]

    for key in required_config_keys:
        if key not in run.config:
            return False

    uid = run.config["uid"]
    validator_hotkey = run.config["hotkey"]
    contest_name = run.config["contest"]

    signing_message = f"{uid}:{validator_hotkey}:{contest_name}"

    try:
        return Keypair(validator_hotkey).verify(signing_message, run.config["signature"])
    except Exception:
        return False


def calculate_score(model_average: float, similarity: float) -> float:
    return max(
        0.0,
        BASELINE_AVERAGE - model_average
    ) * similarity


def get_graph_entries(runs: list[Run]) -> dict[int, GraphEntry]:
    entries: dict[int, GraphEntry] = {}

    for run in reversed(runs[:GRAPH_HISTORY_DAYS]):
        date = datetime.strptime(run.created_at, "%Y-%m-%dT%H:%M:%S")

        for key, value in run.summary.items():
            if key.startswith("_"):
                continue

            if "score" in value:
                continue

            uid = int(key)

            generation_time = value["generation_time"]
            similarity = min(1, value["similarity"])
            score = calculate_score(generation_time, similarity)
            model = value["model"]

            if uid not in entries:
                entries[uid] = GraphEntry([date], [generation_time], [similarity], [score], [model], generation_time)
            else:
                if generation_time < entries[uid].best_time:
                    entries[uid].best_time = generation_time

                data = entries[uid]
                data.dates.append(date)
                data.generation_times.append(data.best_time)
                data.similarities.append(similarity)
                data.scores.append(score)
                data.models.append(model)

    return dict(sorted(entries.items(), key=lambda entry: entry[1].best_time)[:MAX_GRAPH_ENTRIES])


def create_graph(runs: list[Run]) -> go.Figure:
    entries = get_graph_entries(runs)
    fig = go.Figure()

    for uid, data in entries.items():
        fig.add_trace(go.Scatter(
            x=data.dates,
            y=data.generation_times,
            customdata=np.stack((data.similarities, data.scores, data.models), axis=-1),
            mode="lines+markers",
            name=uid,
            hovertemplate=(
                    "<b>Date:</b> %{x|%Y-%m-%d}<br>" +
                    "<b>Generation Time:</b> %{y}s<br>" +
                    "<b>Similarity:</b> %{customdata[0]}<br>" +
                    "<b>Score:</b> %{customdata[1]}<br>" +
                    "<b>Model:</b> %{customdata[2]}<br>"
            ),
        ))

    date_range = max(entries.values(), key=lambda entry: len(entry.dates)).dates

    fig.add_trace(go.Scatter(
        x=date_range,
        y=[BASELINE_AVERAGE] * len(date_range),
        line=dict(color="#ff0000", width=3),
        mode="lines",
        name="Baseline",
    ))

    background_color = gr.themes.default.colors.slate.c800

    fig.update_layout(
        title="Generation Time Improvements",
        yaxis_title="Generation Time (s)",
        plot_bgcolor=background_color,
        paper_bgcolor=background_color,
        template="plotly_dark"
    )

    return fig


def create_leaderboard(runs: list[Run]) -> list[tuple]:
    entries: dict[int, LeaderboardEntry] = {}

    for run in runs:
        has_data = False
        for key, value in run.summary.items():
            if key.startswith("_"):
                continue

            has_data = True

            try:
                uid = int(key)
                generation_time = value.get("generation_time", 0)
                similarity = min(1, value.get("similarity", 0))
                score = value.get("score", calculate_score(generation_time, similarity))

                entries[uid] = LeaderboardEntry(
                    uid=uid,
                    rank=value["rank"],
                    model=value["model"],
                    score=score,
                    model_average=generation_time,
                    similarity=similarity,
                    hotkey=value["hotkey"],
                    previous_day_winner=value["multiday_winner"],
                )
            except Exception:
                continue

        if has_data:
            break

    leaderboard: list[tuple] = [
        (entry.rank + 1, entry.uid, entry.model, entry.score, f"{entry.model_average:.3f}s", f"{entry.similarity:.3f}",
         entry.hotkey, entry.previous_day_winner)
        for entry in sorted(entries.values(), key=lambda entry: (entry.score, entry.rank), reverse=True)
    ]

    return leaderboard


def get_run_validator_uid(run: Run) -> int:
    json_config = json.loads(run.json_config)
    uid = int(json_config["uid"]["value"])
    return uid


def fetch_wandb_data():
    wandb_runs = wandb_api.runs(
        WANDB_RUN_PATH,
        filters={"config.type": "validator"},
        order="-created_at",
    )

    global runs
    runs.clear()
    for run in wandb_runs:
        if not is_valid_run(run):
            continue

        uid = get_run_validator_uid(run)
        if uid not in runs:
            runs[uid] = []
        runs[uid].append(run)

    runs = dict(sorted(runs.items(), key=lambda item: item[0]))


def refresh():
    fetch_wandb_data()
    demo.clear()
    with demo:
        gr.Image(
            "cover.png",
            show_label=False,
            show_download_button=False,
            interactive=False,
            show_fullscreen_button=False,
            show_share_button=False,
        )

        gr.Markdown(
            """
            <center>
            <h1 style="font-size: 50px"> SN39 EdgeMaxxing Leaderboard </h1>
            
            This leaderboard for SN39 tracks the results and top model submissions from current and previous contests.
            </center>
            """)

        with gr.Accordion("Contest #1 Submission Leader: New Dream SDXL on NVIDIA RTX 4090s"):
            choices = list(runs.keys())
            dropdown = gr.Dropdown(
                choices,
                value=SOURCE_VALIDATOR_UID,
                interactive=True,
                label="Source Validator"
            )

            graph = gr.Plot()

            leaderboard = gr.components.Dataframe(
                create_leaderboard(runs[dropdown.value]),
                headers=["Rank", "Uid", "Model", "Score", "Generation Time", "Similarity", "Hotkey",
                         "Previous day winner"],
                datatype=["number", "number", "markdown", "number", "markdown", "markdown", "markdown", "bool"],
                elem_id="leaderboard-table",
                interactive=False,
                visible=True,
            )

            demo.load(lambda uid: create_graph(runs[uid]), [dropdown], [graph])

            dropdown.change(lambda uid: create_graph(runs[uid]), [dropdown], [graph])
            dropdown.change(lambda uid: create_leaderboard(runs[uid]), [dropdown], [leaderboard])


if __name__ == "__main__":
    refresh()
    demo.launch(prevent_thread_lock=True)

    while True:
        time.sleep(REFRESH_RATE)

        now = datetime.now(tz=ZoneInfo("America/New_York"))
        print(f"Refreshing Leaderboard at {now.strftime('%Y-%m-%d %H:%M:%S')}")

        refresh()