import os
from dataclasses import dataclass

from substrateinterface import Keypair
from wandb.apis.public import Runs, Run

import wandb
import gradio as gr

wandb_api = wandb.Api()

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


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


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 main():
    demo = gr.Blocks(css=".typewriter {font-family: 'JMH Typewriter', sans-serif;}")
    with demo:
        with gr.Accordion("Submission Leaderboard"):
            runs: Runs = wandb_api.runs(
                WANDB_RUN_PATH,
                filters={"config.type": "validator", "config.uid": SOURCE_VALIDATOR_UID},
                order="-created_at",
            )

            entries: dict[int, LeaderboardEntry] = {}

            for run in runs:
                if not is_valid_run(run):
                    continue

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

                    try:
                        uid = int(key)

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

                leaderboard: list[tuple] = [
                    (entry.uid, entry.model, entry.score, entry.hotkey, entry.previous_day_winner)
                    for entry in sorted(entries.values(), key=lambda entry: entry.rank)
                ]

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

    demo.launch()


main()