from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta
from typing import TypeAlias

import bittensor as bt

from wandb_data import Uid, TIMEZONE

Weight: TypeAlias = float
Incentive: TypeAlias = float

NET_UID = 39
WEIGHTS_BY_MINER: list[list[tuple[Uid, Weight]]] = []
INCENTIVES: dict[Uid, Incentive] = {}
VALIDATOR_IDENTITIES: dict[Uid, str] = {}
VTRUST: dict[Uid, float] = {}
UPDATED: dict[Uid, int] = {}

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


def fetch_weights():
    WEIGHTS_BY_MINER.clear()
    weights = subtensor.weights(netuid=NET_UID)

    for miner_uid in range(metagraph.n.item()):
        WEIGHTS_BY_MINER.append([])
        for validator_uid, validator_weights in weights:
            if not metagraph.validator_permit[validator_uid]:
                continue

            weight = 0.0
            for miner_weight in validator_weights:
                if miner_weight[0] == miner_uid:
                    weight = miner_weight[1] / 2 ** 16
                    break
            WEIGHTS_BY_MINER[miner_uid].append((validator_uid, weight))


def fetch_incentives():
    INCENTIVES.clear()
    for i in range(len(metagraph.incentive)):
        INCENTIVES[i] = metagraph.incentive[i]


def fetch_vtrust():
    VTRUST.clear()
    for i in range(len(metagraph.validator_trust)):
        VTRUST[i] = metagraph.validator_trust[i]


def fetch_updated(block: int):
    UPDATED.clear()
    for i in range(len(metagraph.last_update)):
        UPDATED[i] = block - metagraph.last_update[i]


def fetch_identities():
    VALIDATOR_IDENTITIES.clear()
    for uid in range(metagraph.n.item()):
        if not metagraph.validator_permit[uid]:
            continue
        identity = subtensor.substrate.query('SubtensorModule', 'Identities', [metagraph.coldkeys[uid]])
        if identity != None:
            VALIDATOR_IDENTITIES[uid] = identity.value["name"]


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


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

    def sync_task():
        print("Syncing metagraph...")
        block = subtensor.get_current_block()
        metagraph.sync(block=block, subtensor=subtensor)

        fetch_weights()
        fetch_incentives()
        fetch_vtrust()
        fetch_updated(block)

        global last_identity_sync
        if now - last_identity_sync < timedelta(hours=12):
            return

        last_identity_sync = now
        fetch_identities()

    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(sync_task)
        try:
            future.result(timeout=timeout)
        except TimeoutError:
            print("Timed out while syncing metagraph")
        except Exception as e:
            print(f"Error occurred while syncing metagraph: {e}")