import gradio as gr
import requests
import time
from threading import Thread

# URLs for the APIs
client1_url = 'https://orionai-training-data-collection-2.hf.space/api/update_token_display'
client2_url = 'https://orionai-training-data-collection-3.hf.space/api/update_token_display'
client3_url = 'https://orionai-training-data-collection.hf.space/api/update_token_display'

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

def get_token_count(url):
    try:
        response = requests.post(url, headers={'Content-Type': 'application/json'})
        response.raise_for_status()  # Raise an error for bad status codes
        result = response.json()
        return int(result.get('token_count', 0))  # Adjust according to API response
    except (requests.RequestException, ValueError) as e:
        print(f"Error fetching token count from {url}: {e}")
        return 0

def monitor_growth():
    while True:
        try:
            curr_count1 = get_token_count(client1_url)
            curr_count2 = get_token_count(client2_url)
            curr_count3 = get_token_count(client3_url)
            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"]

# 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 update_dashboard():
        tokens, speed = get_dashboard_metrics()
        return tokens, speed
    
    # Automatically update the metrics every 5 seconds
    gr.Timer(5000, update_dashboard, [total_tokens, growth_speed], every=5000)

demo.launch()