import os
import gradio as gr
from huggingface_hub import HfApi

HF_TOKEN = os.getenv("HF_TOKEN")
api = HfApi(token=HF_TOKEN)

def scan_spaces_detailed(username: str):
    """
    Scan all Spaces for a user and display metadata in table form.
    Sort all ZeroGPU spaces first.
    """
    try:
        spaces = api.list_spaces(author=username, token=HF_TOKEN)
        if not spaces:
            return [["No Spaces found", "", "", ""]]

        results = []

        for space in spaces:
            try:
                info = api.space_info(space.id, token=HF_TOKEN)
                status = getattr(info, "status", "unknown")
                tags = getattr(info, "tags", [])

                # Determine hardware type
                runtime = getattr(info, "runtime", None)
                if runtime:
                    requested_hw = getattr(runtime, "requested_hardware", None)
                    if requested_hw and requested_hw.lower().startswith("zero"):
                        hardware = "ZeroGPU"
                    else:
                        hardware = "GPU"
                else:
                    hardware_field = getattr(info, "hardware", None)
                    if hardware_field:
                        hardware = "ZeroGPU" if all("cpu" in h.lower() for h in hardware_field) else "GPU"
                    else:
                        hardware = "ZeroGPU" if "cpu" in tags else "GPU"

                results.append([
                    space.id,
                    status,
                    hardware,
                    ", ".join(tags)
                ])

            except Exception as e:
                results.append([
                    space.id,
                    f"Error retrieving info: {e}",
                    "Unknown",
                    ""
                ])

        # Sort ZeroGPU first
        results.sort(key=lambda x: 0 if x[2] == "ZeroGPU" else 1)

        return results

    except Exception as e:
        return [[f"Error: {e}", "", "", ""]]


# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🔍 Hugging Face Spaces ZeroGPU Scanner (Tabular View)")

    with gr.Row():
        username_input = gr.Textbox(label="Hugging Face Username", value="rahul7star")
        scan_button = gr.Button("Scan Spaces")

    output_table = gr.Dataframe(
        headers=["Space ID", "Status", "Hardware", "Tags"],
        datatype=["str", "str", "str", "str"],
        interactive=False,
        wrap=True
    )

    scan_button.click(scan_spaces_detailed, inputs=username_input, outputs=output_table)

demo.launch()