import gradio as gr
from gradio_client import Client
import time
from threading import Thread

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

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

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

def monitor_growth():
    while True:
        try:
            curr_count1 = get_token_count(client1)
            curr_count2 = get_token_count(client2)
            curr_count3 = get_token_count(client3)
            growth_speed1 = curr_count1 - state["prev_count1"]
            growth_speed2 = curr_count2 - state["prev_count2"]
            growth_speed3 = curr_count3 - state["prev_count3"]
            total_tokens = curr_count1 + curr_count2 + curr_count3
            total_growth_speed = growth_speed1 + growth_speed2 + growth_speed3

            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"] = total_growth_speed

            time.sleep(5)
        except Exception as e:
            print(f"Error in monitor growth: {e}")

# Start the monitor thread
Thread(target=monitor_growth, daemon=True).start()

def get_dashboard_metrics():
    return state["prev_total_tokens"], state["total_growth_speed"]

def update_dashboard():
    tokens, speed = get_dashboard_metrics()
    return tokens, speed

# Create Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Token Count Dashboard")
    
    total_tokens = gr.Number(label="Total Token Count", value=0)
    growth_speed = gr.Number(label="Total Growth Speed", value=0)
    
    def refresh_metrics():
        tokens, speed = update_dashboard()
        total_tokens.update(value=tokens)
        growth_speed.update(value=speed)
    
    update_button = gr.Button("Update Stats")
    update_button.click(refresh_metrics)
    
demo.launch()