import gradio as gr
import pandas as pd
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta
from math import ceil
from typing import TypeAlias, Optional, Tuple

import aiohttp
import asyncio

from fiber.chain.interface import get_substrate
from fiber.chain.metagraph import Metagraph
from fiber.chain.models import Node
from substrateinterface.storage import StorageKey

Weight: TypeAlias = float
Hotkey: TypeAlias = str
Uid: TypeAlias = int

NET_UID = 36
VALIDATOR_IDENTITIES: dict[Hotkey, str] = {}
UPDATED: dict[Hotkey, int] = {}
UIDS_BY_HOTKEY: dict[Hotkey, Uid] = {}
HOTKEYS_BY_UID: dict[Uid, Hotkey] = {}

substrate = get_substrate()
metagraph = Metagraph(substrate, netuid=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()
        raise

def is_validator(node: Node) -> bool:
    if not hasattr(node, 'vtrust') or not hasattr(node, 'ip'):
        return False
    return node.vtrust > 0 and str(node.ip) != "0.0.0.0"

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 is not None:
                VALIDATOR_IDENTITIES[hotkey] = info.value["name"]
                break

async def fetch_validator_stats(session: aiohttp.ClientSession, ip: str) -> Tuple[Optional[int], Optional[str], bool]:
    try:
        # Fetch current step (in bits)
        step_url = f"http://{ip}/step"
        async with session.get(step_url, timeout=5) as resp:
            if resp.status == 200:
                step = await resp.json()
                byte_pos = step // 8
                byte_range = f"bytes={byte_pos}-{byte_pos}"
                print(f"Got step {step}, using range: {byte_range}")
            else:
                print(f"Failed to get step from {ip}: {resp.status}")
                return None, None, False
                
        # Get byte at step position
        bits_url = f"http://{ip}/bits"
        headers = {'Range': byte_range}
        async with session.get(bits_url, headers=headers, timeout=5) as resp:
            if resp.status == 206:
                bytes_data = await resp.read()
                if bytes_data:
                    last_byte = format(bytes_data[0], '08b')
                    print(f"Got byte from {ip}: {last_byte}")
                    return step, last_byte, True
                else:
                    print(f"No byte data from {ip}")
                    return None, None, False
            else:
                print(f"Failed byte request for {ip}: {resp.status}")
                return None, None, False
                
    except Exception as e:
        print(f"Error for {ip}: {str(e)}")
        return None, None, False

async def fetch_all_validators(ips: list[str]) -> dict[str, Tuple[Optional[int], Optional[str], bool]]:
    async with aiohttp.ClientSession() as session:
        tasks = [fetch_validator_stats(session, ip) for ip in ips]
        results = await asyncio.gather(*tasks)
        return dict(zip(ips, results))

def get_validator_data() -> pd.DataFrame:
    results = []
    
    ips = [f"{node.ip}:{node.port}" for hotkey, node in metagraph.nodes.items() 
           if is_validator(node)]
    
    stats = asyncio.run(fetch_all_validators(ips))
    
    for hotkey, node in metagraph.nodes.items():
        if not is_validator(node):
            continue
            
        ip = f"{node.ip}:{node.port}"
        step, last_byte, api_up = stats.get(ip, (None, None, False))
        
        validator_info = {
            'Name': VALIDATOR_IDENTITIES.get(hotkey, 'unnamed'),
            'UID': UIDS_BY_HOTKEY.get(hotkey, -1),
            'VTrust': float(node.vtrust),
            'IP': ip,
            'API': '✅' if api_up else '❌',
            'Step': step if step is not None else 'N/A',
            'Last Byte': last_byte if last_byte is not None else 'N/A',
            'Updated': UPDATED.get(hotkey, 0)
        }
        results.append(validator_info)
    
    df = pd.DataFrame(results)
    df['VTrust'] = df['VTrust'].round(4)
    return df.sort_values('VTrust', ascending=False)

# Sync logic
last_sync: datetime = datetime.fromtimestamp(0)
last_identity_sync: datetime = datetime.fromtimestamp(0)

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

    def sync_task():
        print("Syncing metagraph...")
        block = substrate.get_block_number(None)
        metagraph.sync_nodes()

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

        fetch_updated(block)

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

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

# Create Gradio interface
custom_css = """
    :root {
        --primary: #2A4365;          /* Base blue */
        --primary-light: #3B5C8F;    /* Lighter blue for hover states */
        --primary-dark: #1A2F4C;     /* Darker blue for backgrounds */
        --primary-fade: #4A6285;     /* Faded blue for secondary elements */
        --text-primary: #EDF2F7;     /* Light gray for main text */
        --text-secondary: #CBD5E0;   /* Darker gray for secondary text */
    }

    .gradio-container {
        background-image: url('https://huggingface.co/spaces/wombo/pyramid-scheme-validator-states/resolve/main/assets/background2.png') !important;
        background-size: cover !important;
        background-position: center !important;
        background-attachment: fixed !important;
        min-height: 100vh !important;
    }

    /* Style the table */
    .table-wrap {
        background-color: var(--primary) !important;
        border-radius: 8px !important;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1) !important;
        opacity: 0.95 !important;  /* Slight transparency to show background */
    }

    /* Style the table headers */
    thead tr th {
        background-color: var(--primary-fade) !important;
        color: var(--text-primary) !important;
        font-weight: 600 !important;
        border-bottom: 2px solid var(--primary-light) !important;
    }

    /* Style table cells */
    tbody tr td {
        color: var(--text-primary) !important;
        border-bottom: 1px solid var(--primary-light) !important;
    }

    /* Alternating row colors */
    tbody tr:nth-child(even) {
        background-color: var(--primary-dark) !important;
    }

    /* Row hover effect */
    tbody tr:hover {
        background-color: var(--primary-light) !important;
    }

    /* Style the refresh button */
    button.primary {
        background-color: var(--primary) !important;
        color: var(--text-primary) !important;
        border: 1px solid var(--primary-light) !important;
        transition: all 0.2s ease !important;
    }

    button.primary:hover {
        background-color: var(--primary-light) !important;
        transform: translateY(-1px) !important;
    }

    /* Style the title */
    h1 {
        color: var(--text-primary) !important;
        text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.3);
    }

    #title {
        width: 100%;
        text-align: center;
    }
    
    #title h1 {
        width: 100%;
        text-align: center;
        font-size: 3em;
        margin: 0.5em 0;
        display: block;
    }
"""

with gr.Blocks(title="📊 SN36 Validator States", css=custom_css) as demo:
    gr.Markdown("""
    <h1 style="text-align: center; font-size: 3em; margin: 0.5em 0; width: 100%; display: block;">
        📊 SN36 Validator States
    </h1>
    """, elem_id="title")
    
    # Initialize table with empty DataFrame
    sync_metagraph()  # Initial sync
    initial_df = get_validator_data()
    
    table = gr.DataFrame(
        value=initial_df,
        headers=['Name', 'UID', 'VTrust', 'IP', 'API', 'Step', 'Last Byte', 'Updated'],
        datatype=['str', 'number', 'number', 'str', 'str', 'str', 'str', 'number'],
        interactive=False
    )
    
    refresh_btn = gr.Button("Refresh")
    
    def update():
        sync_metagraph()
        return get_validator_data()
    
    refresh_btn.click(fn=update, outputs=table)
    demo.load(fn=update, outputs=table)

if __name__ == "__main__":
    demo.launch(share=False, debug=True, allowed_paths=["assets"])