from pathlib import Path
import gradio as gr

# Define the directory to scan
DATA_PATH = Path("/data")

# Function to get file details and total storage usage
def get_storage():
    # Gather file details
    files = [
        {
            "orig_name": file.name,
            "name": str(file.resolve()),
            "size": f"{file.stat().st_size / (1024.0 ** 2):.2f} MB",  # Size in MB
        }
        for file in DATA_PATH.glob("**/*") if file.is_file()
    ]
    # Calculate total storage usage
    total_size = sum(file.stat().st_size for file in DATA_PATH.glob("**/*") if file.is_file())
    usage = f"{total_size / (1024.0 ** 3):.3f} GB"  # Convert to GB
    return files, usage

# Define the Gradio interface
with gr.Blocks() as app:
    # Title and description
    gr.Markdown("# 📁 File Storage Viewer\nView files and calculate storage usage in the `/data` directory.")

    with gr.Row():
        with gr.Column():
            # Button to trigger file listing
            btn = gr.Button("Fetch Files")

        with gr.Column():
            # Outputs: File list and total storage usage
            files = gr.Dataframe(
                headers=["Original Name", "Path", "Size"],
                interactive=False,
                label="Files",
            )
            storage = gr.Textbox(label="Total Storage Usage", interactive=False)

    # Click action to trigger the `get_storage` function
    btn.click(get_storage, inputs=None, outputs=[files, storage])

# Launch the Gradio app with access to the `/data` directory
app.launch(allowed_paths=["/data"])