import gradio as gr
import os
import json
import random

# Directory to store submissions
DATA_DIR = "submissions"
os.makedirs(DATA_DIR, exist_ok=True)

# Predefined task types
TASK_TYPES = ["Classification", "Regression", "Translation"]

# Colors for task cards
CARD_COLORS = ["#FFDDC1", "#FFABAB", "#FFC3A0", "#D5AAFF", "#85E3FF", "#B9FBC0"]

# Function to handle task submission
def submit_task(task_type, description, yaml_text):
    if not yaml_text.strip():
        return "YAML/Text input cannot be empty."

    # Prepare data
    data = {
        "task_type": task_type,
        "description": description,
        "yaml": yaml_text
    }
    file_path = os.path.join(DATA_DIR, f"{task_type}_{len(os.listdir(DATA_DIR))}.json")
    
    try:
        with open(file_path, "w") as f:
            json.dump(data, f)
        print(f"Saved file: {file_path}, Contents: {data}")
        return f"Task submitted successfully under type '{task_type}'!"
    except Exception as e:
        return f"Error saving task: {e}"

# Function to get tasks by type
def get_tasks_by_type(task_type):
    tasks = []
    for file in os.listdir(DATA_DIR):
        # Skip non-JSON files
        if not file.endswith(".json"):
            continue
        
        try:
            with open(os.path.join(DATA_DIR, file), "r") as f:
                data = json.load(f)
                print(f"File: {file}, Content: {data}")
                # Filter by task type
                if data.get("task_type") == task_type:
                    tasks.append(data)
        except (json.JSONDecodeError, KeyError) as e:
            print(f"Error reading file {file}: {e}")
    return tasks

# Function to dynamically add a new task type
def add_new_task_type(new_type):
    if new_type and new_type not in TASK_TYPES:
        TASK_TYPES.append(new_type)
        return gr.update(choices=TASK_TYPES), f"Task type '{new_type}' added successfully!"
    return gr.update(choices=TASK_TYPES), "Task type already exists or invalid input."

# Function to display tasks as clickable cards
def display_tasks(task_type):
    tasks = get_tasks_by_type(task_type)
    html_content = "<div style='display: flex; flex-wrap: wrap; gap: 10px; color: #000;'>"
    for idx, task in enumerate(tasks):
        color = random.choice(CARD_COLORS)
        html_content += f"""
        <div style='background-color: {color}; color: #000; padding: 10px; border-radius: 5px; cursor: pointer;' onclick="document.getElementById('task-details-{idx}').style.display='block';">
            <b>{task['description']}</b>
        </div>
        <div id='task-details-{idx}' style='display: none; margin-top: 10px; border: 1px solid #ccc; background: #fff; padding: 10px; border-radius: 5px;'>
            <b>Task Type:</b> {task['task_type']}<br>
            <b>Description:</b> {task['description']}<br>
            <b>YAML/Text:</b><pre>{task['yaml']}</pre>
            <button style='margin-top: 10px; color: #000;' onclick="document.getElementById('task-details-{idx}').style.display='none';">Close</button>
            <span style='cursor: pointer; float: right; font-size: 18px; color: #000;' onclick="document.getElementById('task-details-{idx}').style.display='none';">&times;</span>
        </div>
        """
    html_content += "</div>"
    return html_content

# Gradio App
with gr.Blocks() as app:
    gr.Markdown("# Task Configuration Sharing Space")

    with gr.Tabs() as tabs:
        with gr.Tab("Submit Task"):
            gr.Markdown("## Submit a New Task Configuration")
            
            # Input fields for the task submission
            task_type_input = gr.Dropdown(
                label="Task Type",
                choices=TASK_TYPES,
                value="Classification",
                interactive=True
            )
            new_task_input = gr.Textbox(
                label="Add New Task Type",
                placeholder="Enter a new task type",
                interactive=True
            )
            add_task_button = gr.Button("Add Task Type")
            add_task_status = gr.Textbox(label="Status", interactive=False)

            description_input = gr.Textbox(
                label="Task Description",
                placeholder="Provide a brief description of the task.",
                lines=3
            )
            yaml_input = gr.Textbox(
                label="YAML/Text Input",
                placeholder="Paste your YAML or text configuration here.",
                lines=20
            )
            submit_button = gr.Button("Submit Task")
            go_to_view_tab = gr.Button("Go to View Tasks")
            submission_status = gr.Textbox(label="Status", interactive=False)

            # Handle adding new task type
            add_task_button.click(
                add_new_task_type,
                inputs=[new_task_input],
                outputs=[task_type_input, add_task_status]  # Update dropdown and status
            )

            # Handle task submission
            submit_button.click(
                submit_task,
                inputs=[task_type_input, description_input, yaml_input],
                outputs=[submission_status]
            )

            # Button to switch to "View Tasks" tab
            go_to_view_tab.click(
                lambda: gr.Tabs.update(visible_tab="View Tasks"),
                inputs=None,
                outputs=[tabs]
            )

        with gr.Tab("View Tasks"):
            gr.Markdown("## View Submitted Tasks")

            task_type_filter = gr.Dropdown(
                label="Filter by Task Type",
                choices=TASK_TYPES,
                value="Classification",
                interactive=True
            )
            view_button = gr.Button("View Tasks")
            task_display = gr.HTML(label="Submitted Tasks")
            go_to_submit_tab = gr.Button("Go to Submit Task")

            # Handle task display
            view_button.click(
                display_tasks,
                inputs=[task_type_filter],
                outputs=[task_display]
            )

            # Button to switch to "Submit Task" tab
            go_to_submit_tab.click(
                lambda: gr.Tabs.update(visible_tab="Submit Task"),
                inputs=None,
                outputs=[tabs]
            )

app.launch()