from collections import defaultdict
from enum import Enum
from enum import IntEnum

import gradio as gr
import pandas as pd
import plotly.graph_objects as go

from chain_data import sync_chain, fetch_commitments
from src import Key
from src.chain_data import get_neurons
from wandb_data import get_current_runs, Run, get_blacklisted_keys


class DuplicateStatus(IntEnum):
    SAFE = 0
    DUPLICATE = 1
    PENDING_VERIFICATION = 2


class SubmissionStatus(Enum):
    BLACKLISTED = ("Blacklisted", "gray")
    DUPLICATE = ("Duplicate", "gray")
    PENDING = ("Pending", "orange")
    VERIFYING = ("Verifying", "aqua")
    DONE = ("Done", "springgreen")
    INVALID = ("Invalid", "red")

    @staticmethod
    def get_status(run: Run, hotkey: Key, coldkey: Key, block: int, revision: str) -> "SubmissionStatus":
        if is_blacklisted(hotkey, coldkey):
            return SubmissionStatus.BLACKLISTED

        duplicate_status = is_duplicate_submission(hotkey, revision)
        if duplicate_status == DuplicateStatus.DUPLICATE:
            return SubmissionStatus.DUPLICATE
        elif duplicate_status == DuplicateStatus.PENDING_VERIFICATION:
            return SubmissionStatus.VERIFYING

        if hotkey in run.submissions and block > run.submissions[hotkey].info.block and hotkey not in run.invalid_submissions:
            return SubmissionStatus.PENDING

        if hotkey in run.submissions:
            return SubmissionStatus.DONE

        if hotkey in run.invalid_submissions:
            return SubmissionStatus.INVALID

        return SubmissionStatus.PENDING


def is_blacklisted(hotkey: Key, coldkey: Key) -> bool:
    return hotkey in get_blacklisted_keys().hotkeys or coldkey in get_blacklisted_keys().coldkeys


def is_duplicate_submission(hotkey: Key, revision: str) -> DuplicateStatus:
    duplicate_selection = get_blacklisted_keys().duplicate_selection

    if any(
        submission.hotkey == hotkey and submission.revision == revision
        for submission in duplicate_selection.duplicate_submissions
    ):
        return DuplicateStatus.DUPLICATE

    elif any(
        submission.hotkey == hotkey and submission.revision == revision
        for submission in duplicate_selection.safe_submissions
    ):
        return DuplicateStatus.SAFE

    return DuplicateStatus.PENDING_VERIFICATION


DROPDOWN_OPTIONS = [status.value[0] for status in SubmissionStatus]


def create_duplicate_submissions_plot() -> gr.Plot:
    sync_chain()

    submissions_by_coldkey = defaultdict(lambda: ([], []))

    for hotkey, commitment in fetch_commitments().items():
        neuron = get_neurons().get(hotkey)
        if not neuron:
            continue

        coldkey = neuron.coldkey

        if is_blacklisted(hotkey, coldkey):
            continue

        duplicate_status = is_duplicate_submission(hotkey, commitment.revision)
        if duplicate_status == DuplicateStatus.DUPLICATE:
            submissions_by_coldkey[coldkey][0].append(commitment)
        elif duplicate_status == DuplicateStatus.SAFE:
            submissions_by_coldkey[coldkey][1].append(commitment)

    submissions_by_coldkey = dict(sorted(
        submissions_by_coldkey.items(),
        key=lambda x: len(x[1][0]) + len(x[1][1]),
        reverse=True
    ))

    figure = go.Figure()

    for coldkey, (duplicate_commitments, safe_commitments) in submissions_by_coldkey.items():
        urls_safe = [c.url for c in safe_commitments]
        figure.add_trace(go.Bar(
            x=[f"{coldkey[:6]}..."],
            y=[len(safe_commitments)],
            name="Safe",
            marker_color="green",  # type: ignore
            hovertemplate="<br>".join([
                "Submissions (%{y}):<br>" + "<br>".join(urls_safe)
            ]) + "<extra></extra>"
        ))

        urls_duplicate = [c.url for c in duplicate_commitments]
        figure.add_trace(go.Bar(
            x=[f"{coldkey[:6]}..."],
            y=[len(duplicate_commitments)],
            name="Duplicate",
            marker_color="red",  # type: ignore
            hovertemplate="<br>".join([
                "Submissions (%{y}):<br>" + "<br>".join(urls_duplicate)
            ]) + "<extra></extra>"
        ))

    figure.update_layout(
        title="Duplicate Submissions",
        xaxis_title="Coldkey",
        yaxis_title="# of Submissions",
        barmode="stack",
        showlegend=False,
        autosize=True,
        margin=dict(l=50, r=50, t=50, b=50),
        template="plotly_dark"
    )

    return gr.Plot(figure)


def create_submissions(submission_filters: list[str]) -> gr.Dataframe:
    data: list[list] = []
    sync_chain()
    runs = sorted(get_current_runs(), key=lambda run: run.uid)

    for hotkey, commitment in fetch_commitments().items():
        neuron = get_neurons().get(hotkey)
        if not neuron:
            continue
        coldkey = neuron.coldkey

        row = [
            neuron.uid,
            f"[{'/'.join(commitment.url.split('/')[-2:])}]({commitment.url})",
            f"[{commitment.block}](https://taostats.io/block/{commitment.block}/extrinsics)",
            f"[{commitment.revision}]({commitment.url}/commit/{commitment.revision})",
            f"[{hotkey[:6]}...](https://taostats.io/hotkey/{hotkey})",
            f"[{coldkey[:6]}...](https://taostats.io/coldkey/{coldkey})",
            commitment.contest.name,
        ]

        has_data = False
        for run in runs:
            status = SubmissionStatus.get_status(run, hotkey, coldkey, commitment.block, commitment.revision)
            if status.value[0] in submission_filters:
                row.append(f"<span style='color: {status.value[1]}'>{status.value[0]}</span>")
                has_data = True
            else:
                row.append("")

        if not has_data:
            continue

        data.append(row)

    data.sort(key=lambda x: int(x[2].split('[')[1].split(']')[0]), reverse=True)

    columns = ["UID", "Model", "Block", "Revision", "Hotkey", "Coldkey", "Contest"]
    datatype = ["number", "markdown", "markdown", "markdown", "markdown", "markdown", "markdown"]
    for run in runs:
        columns.append(f"{run.uid}")
        datatype.append("markdown")

    return gr.Dataframe(
        pd.DataFrame(data, columns=columns),
        datatype=datatype,
        interactive=False,
        max_height=800,
    )