import gradio as gr
from gradio_client import Client
from collections import deque

# Initialize clients
client1 = Client("orionai/training-data-collection_2")
client2 = Client("orionai/training-data-collection_3")
client3 = Client("orionai/training-data-collection")

# URLs for the clients
urls = {
    client1: "https://orionai-training-data-collection-2.hf.space",
    client2: "https://orionai-training-data-collection-3.hf.space",
    client3: "https://orionai-training-data-collection.hf.space"
}

# Initialize state variables
state = {
    "prev_count1": 0,
    "prev_count2": 0,
    "prev_count3": 0,
    "prev_total_tokens": 0,
    "total_growth_speed": 0
}

token_counts = deque(maxlen=60)  # Store up to 60 seconds of token counts

def get_token_count(client):
    """Fetch the token count from a client."""
    try:
        result = client.predict(api_name="/update_token_display")
        return int(result)
    except Exception as e:
        print(f"Error fetching token count from {client.host}: {e}")
        return 0

def update_dashboard():
    """Update the token count and growth speed."""
    try:
        curr_count1 = get_token_count(client1)
        curr_count2 = get_token_count(client2)
        curr_count3 = get_token_count(client3)
        
        total_tokens = curr_count1 + curr_count2 + curr_count3
        token_counts.append(total_tokens)
        
        if len(token_counts) > 1:
            growth_speed = (token_counts[-1] - token_counts[0]) / len(token_counts)
        else:
            growth_speed = 0
        
        state["prev_count1"] = curr_count1
        state["prev_count2"] = curr_count2
        state["prev_count3"] = curr_count3
        state["prev_total_tokens"] = total_tokens
        state["total_growth_speed"] = growth_speed
        
        counts = {
            client1: curr_count1,
            client2: curr_count2,
            client3: curr_count3
        }
        
        max_client = max(counts, key=counts.get)
        max_tokens = counts[max_client]
        max_url = urls[max_client]

        return total_tokens, growth_speed, max_tokens, max_url
    except Exception as e:
        print(f"Error in update dashboard: {e}")
        return 0, 0, 0, ""

def refresh_metrics():
    """Refresh the metrics for the dashboard."""
    tokens, speed, max_tokens, max_url = update_dashboard()
    return tokens, speed, max_tokens, max_url

# Create Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Token Count Dashboard")
    
    total_tokens = gr.Number(label="Total Token Count")
    growth_speed = gr.Number(label="Average Growth Speed (per second over the last minute)")
    max_tokens = gr.Number(label="Max Token Count from a Single App")
    max_url = gr.HTML(label="Link to App with Max Tokens")
    
    update_button = gr.Button("Update Stats")
    update_button.click(fn=refresh_metrics, inputs=[], outputs=[total_tokens, growth_speed, max_tokens, max_url])
    
demo.launch()