import os
import gradio as gr
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
    return 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
    return 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
    return f"Training model with task: {task}"

# Gradio function to simulate chat-like interface
def gradio_fn(task_input, history):
    task = get_task()

    if task is None:
        # If no task, add a new task and request hardware
        add_task_response = add_task(task_input)
        api.request_space_hardware(repo_id=TRAINING_SPACE_ID, hardware=SpaceHardware.T4_MEDIUM)
        
        # Add the new task response to the chat history
        history.append(("Bot", add_task_response))
        return "", history  # Clear the input box and return updated history
    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_response = train_and_upload(task)
            mark_as_done_response = mark_as_done(task)
            
            # Add responses to history
            history.append(("Bot", train_and_upload_response))
            history.append(("Bot", mark_as_done_response))
            
            # 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)
            history.append(("Bot", "Requesting GPU hardware..."))
        
        return "", history  # Clear the input box and return updated history

# Create the Gradio interface for chat
chat_interface = gr.Interface(
    fn=gradio_fn,
    inputs=[gr.Textbox(label="Enter task name", placeholder="Type your task here...", lines=1)],
    outputs=[gr.Chatbot()],
    live=True,
    title="Task Manager Bot",  # Optional: Title for the interface
    description="Interact with the bot to manage tasks and trigger model training."
)

# Launch the Gradio interface
chat_interface.launch()