import gradio as gr
import os
import time
import shutil
import speedtest
import requests
from pathlib import Path

MOUNT_PATH = "/data"

def check_mount():
    """Verify if the /data storage is actually connected."""
    if os.path.exists(MOUNT_PATH):
        return f"✅ Mount found at {MOUNT_PATH}. Storage Type: {os.path.getsize(MOUNT_PATH) if os.path.isfile(MOUNT_PATH) else 'Directory'}"
    return "❌ Mount point /data not found. Ensure storage is attached in Space Settings."

def test_storage_io(file_size_mb):
    """Measures Write and Read speeds of the persistent storage mount."""
    test_file = os.path.join(MOUNT_PATH, "io_test_temp.bin")
    data = os.urandom(1024 * 1024)  # 1MB of random bytes
    file_size_mb = int(file_size_mb)

    # 1. Write Speed
    start_write = time.time()
    with open(test_file, 'wb') as f:
        for _ in range(file_size_mb):
            f.write(data)
    write_time = time.time() - start_write
    write_speed = file_size_mb / write_time

    # 2. Read Speed
    start_read = time.time()
    with open(test_file, 'rb') as f:
        while f.read(1024 * 1024):
            pass
    read_time = time.time() - start_read
    read_speed = file_size_mb / read_time

    # Cleanup
    if os.path.exists(test_file):
        os.remove(test_file)

    return f"Write: {write_speed:.2f} MB/s | Read: {read_speed:.2f} MB/s"

def test_network_speed():
    """Measures external download/upload speeds using Speedtest."""
    try:
        st = speedtest.Speedtest()
        st.get_best_server()
        download = st.download() / 1_000_000  # Convert to Mbps
        upload = st.upload() / 1_000_000      # Convert to Mbps
        return f"Download: {download:.2f} Mbps | Upload: {upload:.2f} Mbps"
    except Exception as e:
        return f"Error running network test: {str(e)}"

# --- Gradio UI Layout ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"# 🚀 Infrastructure Speed Test (Mount: `{MOUNT_PATH}`)")
    gr.Markdown("Use this to verify if your storage link can handle the 33B model shifting.")
    
    with gr.Row():
        status_btn = gr.Button("🔍 Check Mount Status")
        status_out = gr.Textbox(label="Mount Verification", interactive=False)
        status_btn.click(check_mount, outputs=status_out)

    with gr.Tab("💾 Storage I/O Benchmark"):
        gr.Markdown("Tests the speed of the physical link between your CPU and the /data bucket.")
        size_slider = gr.Slider(minimum=100, maximum=2000, value=500, step=100, label="Test File Size (MB)")
        io_btn = gr.Button("Start I/O Test", variant="primary")
        io_out = gr.Textbox(label="Results", interactive=False)
        io_btn.click(test_storage_io, inputs=size_slider, outputs=io_out)

    with gr.Tab("🌐 Network Throughput"):
        gr.Markdown("Tests the internet speed available to the Space for fetching external model weights.")
        net_btn = gr.Button("Run Speedtest", variant="primary")
        net_out = gr.Textbox(label="Results", interactive=False)
        net_btn.click(test_network_speed, outputs=net_out)

    with gr.Accordion("📂 View Files in /data", open=False):
        def list_files():
            try:
                files = os.listdir(MOUNT_PATH)
                return "\n".join(files) if files else "Empty directory"
            except: return "Cannot access /data"
        file_list = gr.Textbox(label="Current Files")
        refresh_btn = gr.Button("Refresh File List")
        refresh_btn.click(list_files, outputs=file_list)

if __name__ == "__main__":
    demo.launch()