import os
import traceback
from dataclasses import dataclass
from datetime import datetime, timedelta
from enum import Enum
from math import ceil
from typing import TypeAlias

from fiber import constants
from fiber.chain.commitments import _deserialize_commitment_field
from fiber.chain.interface import get_substrate
from fiber.chain.metagraph import Metagraph
from fiber.chain.models import Node
from substrateinterface import SubstrateInterface
from substrateinterface.storage import StorageKey

from network_commitments import Decoder
from src import Key, Uid, TIMEZONE

DISABLE_COMMITMENTS_FETCH = int(os.getenv("DISABLE_COMMITMENTS_FETCH") or 0) > 0

Weight: TypeAlias = float
Incentive: TypeAlias = float


class ContestId(Enum):
    FLUX_NVIDIA_4090 = 0
    SDXL_NEWDREAM_NVIDIA_4090 = 1


@dataclass
class Commitment:
    provider: str
    repository: str
    revision: str
    contest: ContestId
    block: int

    @classmethod
    def decode(cls, decoder: Decoder, block: int):
        provider = decoder.read_str()
        repository = decoder.read_str()
        revision = decoder.read_sized_str(7)
        contest_id = ContestId(decoder.read_uint16())

        return cls(
            provider=provider,
            repository=repository,
            revision=revision,
            contest=contest_id,
            block=block
        )

    def get_repo_link(self):
        return f"https://{self.provider}/{self.repository}"


SPEC_VERSION = 8
NET_UID = 39
WEIGHTS_BY_MINER: dict[Key, list[tuple[Key, Weight]]] = {}
VALIDATOR_IDENTITIES: dict[Key, str] = {}
COMMITMENTS: dict[Key, Commitment] = {}
UPDATED: dict[Key, int] = {}

UIDS_BY_HOTKEY: dict[Key, Uid] = {}
HOTKEYS_BY_UID: dict[Uid, Key] = {}

substrate = get_substrate(subtensor_address=constants.FINNEY_SUBTENSOR_ADDRESS)
metagraph = Metagraph(substrate, netuid=str(NET_UID), load_old_nodes=False)


def query_subtensor(storage_keys: list[StorageKey], block: int) -> list:
    global substrate
    try:
        return substrate.query_multi(
            storage_keys=storage_keys,
            block_hash=substrate.get_block_hash(block),
        )
    except Exception:
        substrate = get_substrate(subtensor_address=substrate.url)
        raise


def is_validator(node: Node) -> bool:
    return node.vtrust > 0 or node.stake > 10_000


def get_nodes() -> dict[Key, Node]:
    return metagraph.nodes


def fetch_weights(block: int):
    WEIGHTS_BY_MINER.clear()
    storage_keys: list[StorageKey] = []
    for hotkey, node in metagraph.nodes.items():
        if not is_validator(node): continue
        storage_keys.append(substrate.create_storage_key(
            "SubtensorModule",
            "Weights",
            [metagraph.netuid, UIDS_BY_HOTKEY[node.hotkey]]
        ))

    weights = query_subtensor(storage_keys, block)

    for hotkey, node in metagraph.nodes.items():
        for storage, validator_weights in weights:
            validator_hotkey = HOTKEYS_BY_UID[storage.params[1]]
            if hotkey not in WEIGHTS_BY_MINER:
                WEIGHTS_BY_MINER[hotkey] = []
            weight = 0.0
            for miner_weight in validator_weights:
                if miner_weight[0].value == UIDS_BY_HOTKEY[hotkey]:
                    weight = miner_weight[1].value / 2 ** 16
                    break
            WEIGHTS_BY_MINER[hotkey].append((validator_hotkey, weight))


def fetch_updated(block: int):
    UPDATED.clear()
    for hotkey, node in metagraph.nodes.items():
        UPDATED[hotkey] = ceil(block - node.last_updated)


def fetch_identities(block: int):
    VALIDATOR_IDENTITIES.clear()
    storage_keys: list[StorageKey] = []
    for hotkey, node in metagraph.nodes.items():
        if not is_validator(node): continue
        storage_keys.append(substrate.create_storage_key(
            "SubtensorModule",
            "Identities",
            [node.coldkey]
        ))

    identities = query_subtensor(storage_keys, block)
    for hotkey, node in metagraph.nodes.items():
        for storage, info in identities:
            if node.coldkey != storage.params[0]: continue
            if info != None:  # noqa
                VALIDATOR_IDENTITIES[hotkey] = info.value["name"]
                break


def fetch_commitments(block: int):
    if DISABLE_COMMITMENTS_FETCH:
        return

    COMMITMENTS.clear()
    storage_keys: list[StorageKey] = []
    for hotkey, node in metagraph.nodes.items():
        if is_validator(node): continue
        storage_keys.append(substrate.create_storage_key(
            "Commitments",
            "CommitmentOf",
            [metagraph.netuid, hotkey]
        ))

    commitments = query_subtensor(storage_keys, block)
    for storage, commitment in commitments:
        try:
            if not commitment or not commitment.value:
                continue

            fields = commitment.value["info"]["fields"]
            if not fields:
                continue

            field = _deserialize_commitment_field(fields[0])
            if field is None:
                continue

            decoder = Decoder(field[1])
            spec_version = decoder.read_uint16()
            if spec_version != SPEC_VERSION:
                continue

            COMMITMENTS[storage.params[1]] = Commitment.decode(decoder, int(commitment.value["block"]))
        except:
            continue


last_sync: datetime = datetime.fromtimestamp(0, TIMEZONE)
last_identity_sync: datetime = datetime.fromtimestamp(0, TIMEZONE)
last_commitment_sync: datetime = datetime.fromtimestamp(0, TIMEZONE)


def sync_metagraph():
    global substrate
    global last_sync
    now = datetime.now(TIMEZONE)
    if now - last_sync < timedelta(minutes=5):
        return
    last_sync = now

    try:
        print("Syncing metagraph...")
        metagraph.sync_nodes()
        block = substrate.get_block_number(None)  # type: ignore

        for uid, node in enumerate(metagraph.nodes.values()):
            UIDS_BY_HOTKEY[node.hotkey] = uid
            HOTKEYS_BY_UID[uid] = node.hotkey

        fetch_weights(block)
        fetch_updated(block)

        global last_identity_sync
        if now - last_identity_sync > timedelta(days=1):
            print("Syncing identities...")
            last_identity_sync = now
            fetch_identities(block)

        global last_commitment_sync
        if now - last_commitment_sync > timedelta(hours=12):
            print("Syncing commitments...")
            last_commitment_sync = now
            fetch_commitments(block)
    except Exception:
        print(f"Error occurred while syncing metagraph")
        traceback.print_exc()
        substrate = SubstrateInterface(substrate.url)