import os
import time
import math
import shutil
import platform
import pandas as pd
import gradio as gr

# DEFAULT_DIR = "/bucket_data" 

def run_benchmark(storage_choice, file_size_mb, block_size_kb):
    # Ensure directory exists
    if storage_choice == "Local Data Space (/local_data)":
        target_dir = "/local_data"
    elif storage_choice == "Same Acc Bucket Data Space (/sameAcc_bucket_data)":
        target_dir = "/sameAcc_bucket_data"
    elif storage_choice == "Different Acc Bucket Data Space (/diffAcc_bucket_data)":
        target_dir = "/diffAcc_bucket_data"
    else:
        target_dir = "/local_data"
    
    os.makedirs(target_dir, exist_ok=True)
    test_file_path = os.path.join(target_dir, "__io_perf_test__.bin")
    
    # Conversions
    file_size_bytes = int(file_size_mb * 1024 * 1024)
    block_size_bytes = int(block_size_kb * 1024)
    num_blocks = math.ceil(file_size_bytes / block_size_bytes)
    
    # Pre-generate a data block to avoid processing overhead during I/O loop
    dummy_block = os.urandom(block_size_bytes)
    
    results = {}
    
    try:
        # --- WRITE TEST ---
        # Attempt O_DIRECT to bypass OS caching if supported by the OS/FileSystem
        flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
        if hasattr(os, 'O_DIRECT') and platform.system() == 'Linux':
            try:
                # O_DIRECT requires block sizes to be aligned (usually 4096 bytes)
                fd = os.open(test_file_path, flags | os.O_DIRECT)
            except OSError:
                fd = os.open(test_file_path, flags)  # Fallback
        else:
            fd = os.open(test_file_path, flags)
            
        start_time = time.perf_counter()
        for _ in range(num_blocks):
            os.write(fd, dummy_block)
            
        # Force write out to physical disk before stopping timer
        try:
            os.fdatasync(fd)
        except AttributeError:
            os.fsync(fd)
            
        end_time = time.perf_counter()
        os.close(fd)
        
        write_duration = end_time - start_time
        write_speed = file_size_mb / write_duration
        results['Write Speed'] = f"{write_speed:.2f} MB/s"
        results['Write Time'] = f"{write_duration:.4f} seconds"
        
        # --- READ TEST ---
        flags = os.O_RDONLY
        if hasattr(os, 'O_DIRECT') and platform.system() == 'Linux':
            try:
                fd = os.open(test_file_path, flags | os.O_DIRECT)
            except OSError:
                fd = os.open(test_file_path, flags)
        else:
            fd = os.open(test_file_path, flags)
            
        start_time = time.perf_counter()
        while True:
            read_data = os.read(fd, block_size_bytes)
            if not read_data:
                break
        end_time = time.perf_counter()
        os.close(fd)
        
        read_duration = end_time - start_time
        read_speed = file_size_mb / read_duration
        results['Read Speed'] = f"{read_speed:.2f} MB/s"
        results['Read Time'] = f"{read_duration:.4f} seconds"
        
        # --- SYSTEM INFO ---
        total, used, free = shutil.disk_usage(target_dir)
        results['Storage Total'] = f"{total / (1024**3):.2f} GB"
        results['Storage Free'] = f"{free / (1024**3):.2f} GB"
        results['Status'] = "Success"
        
    except Exception as e:
        results['Status'] = f"Error: {str(e)}"
        results['Write Speed'], results['Read Speed'] = "N/A", "N/A"
        results['Write Time'], results['Read Time'] = "N/A", "N/A"
        results['Storage Total'], results['Storage Free'] = "N/A", "N/A"
    finally:
        # Cleanup test file safely
        if os.path.exists(test_file_path):
            os.remove(test_file_path)
            
    # Format to Pandas DataFrame for a pretty UI presentation
    df_data = {
        "Metric": list(results.keys()),
        "Value": list(results.values())
    }
    return pd.DataFrame(df_data)


# --- GRADIO INTERFACE ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 💾 Hugging Face Space local /data Speed Test
        Use this utility to benchmark the real-world read/write performance of ephemeral paths or attached persistent storage (`/data`).
        """
    )
    
    with gr.Row():
        with gr.Column():
            storage_choice = gr.Radio(
                label="Select Target Storage Volume",
                choices=["Local Data Space (/local_data)", "Same Acc Bucket Data Space (/sameAcc_bucket_data)", "Different Acc Bucket Data Space (/diffAcc_bucket_data)"],
                value="Local Data Space (/local_data)",
                info="Switches underlying workspace mounts directly."
            )               
            # test_path = gr.Textbox(
            #     label="Target Directory", 
            #     value=DEFAULT_DIR, 
            #     placeholder="e.g. /data or ./test_dir"
            # )
            size_slider = gr.Slider(
                label="Test File Size (MB)", 
                minimum=10, 
                maximum=10000, 
                value=100, 
                step=10
            )
            block_slider = gr.Slider(
                label="Block Size (KB) - UI Buffer sizing", 
                minimum=4, 
                maximum=4096, 
                value=1024, 
                step=4
            )
            btn_run = gr.Button("🚀 Run IO Benchmark", variant="primary")
            
        with gr.Column():
            gr.Markdown("### 📊 Benchmark Results")
            output_table = gr.DataFrame(
                headers=["Metric", "Value"],
                datatype=["str", "str"],
                interactive=False
            )

    btn_run.click(
        fn=run_benchmark,
        inputs=[storage_choice, size_slider, block_slider],
        outputs=output_table
    )

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