import os
import streamlit as st
from huggingface_hub import HfApi, SpaceHardware

# Set up Hugging Face API token and Space ID
HF_TOKEN = os.getenv("HF_TOKEN")  # Ensure your Hugging Face token is set as a secret
TRAINING_SPACE_ID = "your_space_id_here"  # Replace with your actual space ID

# Initialize Hugging Face API
api = HfApi(token=HF_TOKEN)

# Function to check for a scheduled task (this is a placeholder for your actual task-checking logic)
def get_task():
    # You can implement logic here to check for scheduled tasks
    return None  # For example, return None if no task is scheduled

# Function to add a new task (you can implement this depending on your use case)
def add_task(task):
    # Logic to add a new task
    st.write(f"Task '{task}' added!")

# Function to mark the task as "DONE" (this is a placeholder)
def mark_as_done(task):
    # Mark the task as done once it's completed
    st.write(f"Task '{task}' completed!")

# Function to simulate training the model (replace with actual training logic)
def train_and_upload(task):
    # Implement your model training logic here
    st.write(f"Training model with task: {task}")

# Check if there’s an existing task
task = get_task()

if task is None:
    # Display Gradio interface to request a new task
    def gradio_fn(task):
        # On user request, add task and request hardware
        add_task(task)
        api.request_space_hardware(repo_id=TRAINING_SPACE_ID, hardware=SpaceHardware.T4_MEDIUM)

    # Use Streamlit to request a task (Gradio interface or a simple button to simulate this)
    task_input = st.text_input("Enter task name", "")
    if st.button("Request Task"):
        gradio_fn(task_input)
else:
    # If a task is available, check for hardware
    runtime = api.get_space_runtime(repo_id=TRAINING_SPACE_ID)
    if runtime.hardware == SpaceHardware.T4_MEDIUM:
        # Fine-tune model on GPU if available
        train_and_upload(task)

        # Mark task as "DONE" after training
        mark_as_done(task)

        # Reset to CPU hardware after training
        api.request_space_hardware(repo_id=TRAINING_SPACE_ID, hardware=SpaceHardware.CPU_BASIC)
    else:
        # If GPU hardware is not available, request it
        api.request_space_hardware(repo_id=TRAINING_SPACE_ID, hardware=SpaceHardware.T4_MEDIUM)