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

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

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

START_DATE = datetime(2024, 9, 17)
NET_UID = 39
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;}", fill_height=True, fill_width=True)

subtensor = bt.subtensor()
metagraph = bt.metagraph(netuid=NET_UID)
bt.logging.disable_logging()

runs: dict[int, list[Run]] = {}
validator_identities: dict[int, str] = {}


@dataclass
class LeaderboardEntry:
    uid: int
    winner: bool
    repository: str
    score: float
    similarity: float
    hotkey: str
    baseline_generation_time: float
    generation_time: float
    size: int
    vram_used: float
    watts_used: float


@dataclass
class GraphEntry:
    dates: list[datetime]
    baseline_generation_times: list[float]
    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(baseline_generation_time: float, generation_time: float, similarity_score: float) -> float:
    return (baseline_generation_time - generation_time) * similarity_score


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:%SZ")

        for summary_key, summary_value in run.summary.items():
            if not summary_key.startswith("benchmarks"):
                continue
            for key, value in summary_value.items():
                if "score" in value:
                    continue

                uid = int(key)

                baseline_generation_time = value["baseline_generation_time"]
                generation_time = value["generation_time"]
                similarity = min(1, value["similarity"])
                score = calculate_score(baseline_generation_time, generation_time, similarity)
                model = run.summary["submissions"][str(uid)]["repository"]

                if uid not in entries:
                    entries[uid] = GraphEntry([date], [baseline_generation_time], [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.baseline_generation_times.append(baseline_generation_time)
                    data.generation_times.append(data.best_time)
                    data.similarities.append(similarity)
                    data.scores.append(score)
                    data.models.append(model)

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


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

    average_baseline_generation_times = sum(entry.baseline_generation_times[0] for entry in entries.values()) / len(entries)
    fig.add_trace(go.Scatter(
        x=date_range,
        y=[average_baseline_generation_times] * 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 summary_key, summary_value in run.summary.items():
            if not summary_key == "benchmarks":
                continue
            for key, value in summary_value.items():
                has_data = True

                uid = int(key)
                generation_time = value["generation_time"]
                baseline_generation_time = value["baseline_generation_time"]
                similarity = min(1, value["similarity"])

                entries[uid] = LeaderboardEntry(
                    uid=uid,
                    winner="winner" in value,
                    repository=run.summary["submissions"][str(uid)]["repository"],
                    score=calculate_score(baseline_generation_time, generation_time, similarity),
                    similarity=similarity,
                    baseline_generation_time=baseline_generation_time,
                    generation_time=generation_time,
                    size=value["size"],
                    vram_used=value["vram_used"],
                    watts_used=value["watts_used"],
                    hotkey=value["hotkey"],
                )

        if has_data:
            break

    return [(
        entry.uid,
        f"<span style='color: {'springgreen' if entry.winner else 'red'}'>{entry.winner}</span>",
        entry.repository,
        round(entry.score, 3),
        f"{entry.generation_time:.3f}s",
        f"{entry.similarity:.3f}",
        f"{entry.size / 1_000_000_000:.3f}GB",
        f"{entry.vram_used / 1_000_000_000:.3f}GB",
        f"{entry.watts_used:.3f}W",
        entry.hotkey,
    ) for entry in sorted(entries.values(), key=lambda entry: (entry.winner, entry.score), reverse=True)]


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", "created_at": {'$gt': str(START_DATE)}},
        order="-created_at",
    )

    wandb_runs = [run for run in wandb_runs if "benchmarks" in run.summary]

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

        uid = get_run_validator_uid(run)
        if not metagraph.validator_permit[uid]:
            continue

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

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


def fetch_identities():
    validator_identities.clear()
    for uid in runs.keys():
        identity = subtensor.substrate.query('SubtensorModule', 'Identities', [metagraph.coldkeys[uid]])
        if identity != None:
            validator_identities[uid] = identity.value["name"]


def get_validator_name(validator_uid: int) -> str:
    if validator_uid in validator_identities:
        return validator_identities[validator_uid]
    else:
        return metagraph.hotkeys[validator_uid]


def get_choices() -> list[tuple[str, int]]:
    choices: list[tuple[str, int]] = []
    for uid, run in runs.items():
        benchmarks = dict(run[0].summary.get("benchmarks", {}))
        finished = any("winner" in value for value in benchmarks.values())
        progress_text = "Finished" if finished else "In Progress"
        choices.append((f"{uid} - {get_validator_name(uid)} ({progress_text})", uid))
    return choices


def refresh():
    metagraph.sync(subtensor=subtensor)
    fetch_wandb_data()
    fetch_identities()
    demo.clear()
    now = datetime.now(tz=ZoneInfo("America/New_York"))

    with demo:
        gr.Image(
            "cover.png",
            show_label=False,
            show_download_button=False,
            interactive=False,
            show_fullscreen_button=False,
            show_share_button=False,
            container=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(f"Contest #1 Submission Leader: New Dream SDXL on NVIDIA RTX 4090s (Last updated: {now.strftime('%Y-%m-%d %I:%M:%S %p')} EST)"):
            dropdown = gr.Dropdown(
                get_choices(),
                value=SOURCE_VALIDATOR_UID,
                interactive=True,
                label="Source Validator"
            )

            leaderboard = gr.components.Dataframe(
                create_leaderboard(runs[dropdown.value]),
                headers=["Uid", "Winner", "Model", "Score", "Gen Time", "Similarity", "Size", "VRAM Usage", "Power Usage", "Hotkey"],
                datatype=["number", "markdown", "markdown", "number", "markdown", "number", "markdown", "markdown", "markdown", "markdown"],
                elem_id="leaderboard-table",
            )

            graph = gr.Plot()
            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()