import gradio as gr
import requests

from huggingface_hub import list_models, list_datasets, list_spaces
from typing import Union

# Helper function to get the total storage for models, datasets, or spaces
def get_total_storage(namespace, resource_type, oauth_token: Union[gr.OAuthToken, None]):
    token = oauth_token.token if oauth_token else None
    if resource_type == "model":
        resources = list(list_models(author=namespace, token=token))
        url_base = "https://huggingface.co/api/models"
    elif resource_type == "dataset":
        resources = list(list_datasets(author=namespace, token=token))
        url_base = "https://huggingface.co/api/datasets"
    elif resource_type == "space":
        resources = list(list_spaces(author=namespace, token=token))
        url_base = "https://huggingface.co/api/spaces"
    
    total_storage = 0
    for resource in resources:
        resource_id = resource.id
        url = f"{url_base}/{resource_id}/treesize/main"
        response = requests.get(url)
        if response.status_code == 200:
            size_info = response.json()
            total_storage += size_info.get("size", 0)

    return total_storage, len(resources)

def get_report(namespace, oauth_token: Union[gr.OAuthToken, None]):
    # Fetch storage and counts for models, datasets, and spaces
    model_storage, n_models = get_total_storage(namespace, "model", oauth_token)
    dataset_storage, n_datasets = get_total_storage(namespace, "dataset", oauth_token)
    space_storage, n_spaces = get_total_storage(namespace, "space", oauth_token)
    
    # Total storage
    total_storage = model_storage + dataset_storage + space_storage
    total_storage_gb = total_storage / (1024 ** 3)  # Convert from bytes to GB
    total_storage_tb = total_storage_gb / 1024       # Convert from GB to TB

    # Cost calculation (1 TB = 20 USD)
    estimated_cost = total_storage_tb * 20
    
    # Generate a report
    report = f"""
    ## Hugging Face Storage Report for {namespace}
    - **Number of Models**: {n_models}
    - **Number of Datasets**: {n_datasets}
    - **Number of Spaces**: {n_spaces}
    
    - **Total Storage**: {total_storage_gb:.2f} GB ({total_storage_tb:.2f} TB)
    - **Estimated Cost**: ${estimated_cost:.2f} USD (at 1 TB = $20USD)
    """
    
    return report

css = """
.main_ui_logged_out{opacity: 0.3; pointer-events: none}
"""

# Create Gradio UI
with gr.Blocks(css=css) as demo:
    gr.Markdown("# Hugging Face Storage Report")
    
    gr.LoginButton()
    namespace = gr.Textbox(label="Enter Namespace (username or org)")
    output = gr.Markdown()
    
    # Button to trigger the report generation
    report_button = gr.Button("Generate Report")
    report_button.click(fn=get_report, inputs=namespace, outputs=output, concurrency_limit=10)

# Launch the Gradio app
demo.launch()