import gradio as gr
import psutil
import os
import time

def get_space_metrics():
    """
    Calculates the memory and CPU usage of the current process 
    and all its children to stay within Space limits.
    """
    try:
        current_process = psutil.Process(os.getpid())
        
        # Calculate total RAM (RSS) for this process and children
        total_rss = current_process.memory_info().rss
        for child in current_process.children(recursive=True):
            total_rss += child.memory_info().rss
        
        ram_usage_mb = total_rss / (1024**2)
        
        # CPU usage over a short interval
        cpu_usage_pct = current_process.cpu_percent(interval=0.1)
        
        return f"RAM: {ram_usage_mb:.1f} MB", f"CPU: {cpu_usage_pct:.1f}%"
    except Exception as e:
        return "Error", str(e)

# Build the Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Space Resource Monitor")
    
    with gr.Row():
        ram_output = gr.Textbox(label="Allocated RAM Usage")
        cpu_output = gr.Textbox(label="Allocated CPU Usage")
    
    refresh_btn = gr.Button("Refresh Metrics")
    refresh_btn.click(get_space_metrics, outputs=[ram_output, cpu_output])
    
    # Auto-refresh on load
    demo.load(get_space_metrics, outputs=[ram_output, cpu_output])

if __name__ == "__main__":
    # CRITICAL: server_name="0.0.0.0" allows external access to the container
    # server_port=7860 is the default port expected by HF Spaces
    demo.launch(server_name="0.0.0.0", server_port=7860)