import asyncio
import json
import os
import signal
from pathlib import Path
from dotenv import load_dotenv
load_dotenv()

def kill_process_on_port(port):
    try:
        # Hex representation of port (e.g. 7860 -> 1EB4)
        port_hex = f"{port:04X}"
        inodes = []
        if os.path.exists("/proc/net/tcp"):
            with open("/proc/net/tcp", "r") as f:
                for line in f:
                    parts = line.split()
                    if len(parts) > 9:
                        local_addr = parts[1]
                        local_port = local_addr.split(":")[-1]
                        if local_port == port_hex:
                            inodes.append(parts[9])
        
        if inodes:
            my_pid = os.getpid()
            for pid_str in os.listdir("/proc"):
                if not pid_str.isdigit():
                    continue
                pid = int(pid_str)
                if pid == my_pid:
                    continue
                fd_dir = f"/proc/{pid_str}/fd"
                if not os.path.exists(fd_dir):
                    continue
                try:
                    for fd in os.listdir(fd_dir):
                        link = os.readlink(f"{fd_dir}/{fd}")
                        for inode in inodes:
                            if f"socket:[{inode}]" in link:
                                print(f"Self-Healing: Killing zombie process {pid} using port {port}", flush=True)
                                os.kill(pid, signal.SIGKILL)
                                break
                except Exception:
                    continue
    except Exception as e:
        print(f"Error checking/killing process on port {port}: {e}", flush=True)

# Free port 7860 before trying to launch Gradio
kill_process_on_port(7860)


import gradio as gr

# --- Gradio 5.x BugFix Monkeypatch ---
import gradio_client.utils as client_utils
_original_json_schema_to_python_type = client_utils._json_schema_to_python_type
def _safe_json_schema_to_python_type(schema, defs=None):
    if isinstance(schema, bool):
        return "Any"
    return _original_json_schema_to_python_type(schema, defs)
client_utils._json_schema_to_python_type = _safe_json_schema_to_python_type
# -------------------------------------
ROOT = Path(__file__).parent
GRAMMAR_PATH = ROOT / "grammar.gbnf"

# Keep the demo runnable while the GGUF files are still being prepared.
MOCK_MODE = os.getenv("STEP_ZERO_MOCK", "1") != "0"
NEMOTRON_MODEL_PATH = os.getenv("NEMOTRON_MODEL_PATH", "./models/step-zero-nemotron-finetuned.gguf")
MINICPM_MODEL_PATH = os.getenv("MINICPM_MODEL_PATH", "./models/minicpm-3-4b.gguf")

if not MOCK_MODE:
    import os
    from huggingface_hub import hf_hub_download
    from llama_cpp import Llama
    from llama_cpp import LlamaGrammar

    os.makedirs("./models", exist_ok=True)
    
    if not os.path.exists(NEMOTRON_MODEL_PATH):
        print(f"Downloading Nemotron model to {NEMOTRON_MODEL_PATH}...", flush=True)
        hf_hub_download(
            repo_id="tc043/step-zero-nemotron",
            filename="step-zero-nemotron-finetuned.gguf",
            local_dir="./models"
        )
        
    if not os.path.exists(MINICPM_MODEL_PATH):
        print(f"Downloading MiniCPM model to {MINICPM_MODEL_PATH}...", flush=True)
        downloaded = hf_hub_download(
            repo_id="mradermacher/MiniCPM3-4B-GGUF",
            filename="MiniCPM3-4B.Q4_K_M.gguf",
            local_dir="./models"
        )
        expected_path = os.path.abspath(MINICPM_MODEL_PATH)
        downloaded_path = os.path.abspath(downloaded)
        if downloaded_path != expected_path and os.path.exists(downloaded_path):
            os.rename(downloaded_path, expected_path)

    grammar = LlamaGrammar.from_file(str(GRAMMAR_PATH))

    # Models start as None and are loaded in background threads immediately
    nemotron = None
    minicpm = None

    import threading

    _nemotron_ready = threading.Event()
    _minicpm_ready = threading.Event()

    def _preload_nemotron():
        global nemotron
        print(f"BACKGROUND PRELOAD: Loading Nemotron from {NEMOTRON_MODEL_PATH}...", flush=True)
        nemotron = Llama(model_path=NEMOTRON_MODEL_PATH, n_ctx=1024)
        _nemotron_ready.set()
        print("BACKGROUND PRELOAD: Nemotron ready.", flush=True)

    def _preload_minicpm():
        global minicpm
        print(f"BACKGROUND PRELOAD: Loading MiniCPM from {MINICPM_MODEL_PATH}...", flush=True)
        minicpm = Llama(model_path=MINICPM_MODEL_PATH, n_ctx=1024)
        _minicpm_ready.set()
        print("BACKGROUND PRELOAD: MiniCPM ready.", flush=True)

    # Fire off both loads in parallel — Gradio server boots independently
    threading.Thread(target=_preload_nemotron, daemon=True).start()
    threading.Thread(target=_preload_minicpm, daemon=True).start()

# Per-model locks for concurrency
nemotron_lock = asyncio.Lock()
minicpm_lock = asyncio.Lock()

async def run_nemotron(prompt, **kwargs):
    async with nemotron_lock:
        if not MOCK_MODE:
            _nemotron_ready.wait()  # blocks only if still loading
        return await asyncio.to_thread(nemotron, prompt, **kwargs)

async def run_minicpm_chat(messages, **kwargs):
    async with minicpm_lock:
        if not MOCK_MODE:
            _minicpm_ready.wait()  # blocks only if still loading
        return await asyncio.to_thread(minicpm.create_chat_completion, messages=messages, **kwargs)

MAX_TRACE_EVENTS = 50
MAX_HISTORY_ITEMS = 20

def append_trace(state, event: dict) -> None:
    state["trace"].append(event)
    if len(state["trace"]) > MAX_TRACE_EVENTS:
        state["trace"] = state["trace"][-MAX_TRACE_EVENTS:]

def append_history(state, displayed_task: str, raw_task: str) -> None:
    state["history"].append(displayed_task)
    state["raw_history"].append(raw_task)
    if len(state["history"]) > MAX_HISTORY_ITEMS:
        state["history"] = state["history"][-MAX_HISTORY_ITEMS:]
    if len(state["raw_history"]) > MAX_HISTORY_ITEMS:
        state["raw_history"] = state["raw_history"][-MAX_HISTORY_ITEMS:]

def is_semantic_repeat(new_task: str, history: list) -> bool:
    if not history:
        return False
        
    stop_words = {"i", "need", "to", "a", "the", "an", "and", "or", "but", "with", "for", "of", "on", "at", "by", "start", "with"}
    
    def get_stems(task_str):
        words = task_str.lower().split()
        stems = []
        for w in words:
            w = w.strip(".,!?\"'();:")
            if w in stop_words or not w:
                continue
            # Simple suffix stemming
            if len(w) > 4:
                if w.endswith("ing"):
                    w = w[:-3]
                elif w.endswith("ies"):
                    w = w[:-3] + "y"
                elif w.endswith("es"):
                    w = w[:-2]
                elif w.endswith("ed"):
                    w = w[:-2]
                elif w.endswith("s") and not w.endswith("ss"):
                    w = w[:-1]
            stems.append(w)
        return set(stems)

    new_words = get_stems(new_task)
    if not new_words:
        return False
        
    for past in history[-3:]:
        past_words = get_stems(past)
        if not past_words:
            continue
        overlap = len(new_words & past_words) / max(len(new_words | past_words), 1)
        if overlap > 0.6:  # 60% overlap on key word stems = semantic repeat
            return True
    return False

def clean_and_validate_task(raw_output: str, history: list, skipped_history: list) -> tuple[str, bool]:
    import re
    if not raw_output:
        return "", False
        
    # Split by sentence boundaries and take the first sentence
    sentences = re.split(r'(?<=[.!?])\s+', raw_output)
    first_sentence = sentences[0].strip() if sentences else raw_output
    
    # Strip quotes, punctuation, numbers, etc. at ends
    cleaned = first_sentence.strip('"\' ,.-1234567890)')
    if not cleaned:
        return "", False
        
    # Strip conversational prefixes
    prefixes_to_strip = [
        "i need to ", "i should ", "i have to ", "let's ", "we should ",
        "please ", "you need to ", "you should "
    ]
    cleaned_lower = cleaned.lower()
    for prefix in prefixes_to_strip:
        if cleaned_lower.startswith(prefix):
            cleaned = cleaned[len(prefix):]
            cleaned_lower = cleaned.lower()
            
    # Capitalize the first letter again
    if cleaned:
        cleaned = cleaned[0].upper() + cleaned[1:]
        
    # Check for prompt leakage and scaffolding words
    invalid_keywords = [
        "completed tasks", "failures", "system", "user", "assistant", 
        "next step", "done", "fail", "gbnf", "pacemaker", "extra_id_",
        "instruction", "output", "goal:", "failures:"
    ]
    contains_scaffold = any(k in cleaned.lower() for k in invalid_keywords)
    
    # Check for reasonable length (under 12 words)
    word_count = len(cleaned.split())
    if word_count > 12:
        return cleaned, False
        
    if contains_scaffold:
        return cleaned, False
        
    if is_semantic_repeat(cleaned, history + skipped_history):
        return cleaned, False
        
    return cleaned, True

async def generate_atomic_task(goal: str, previous_failures: int, history: list = None, skipped_history: list = None, rejected_task: str = None) -> str:
    if history is None:
        history = []
    if skipped_history is None:
        skipped_history = []
    
    recent_history = history[-3:] if history else []
    history_str = "\n".join([f"- {t}" for t in recent_history]) if recent_history else "None"
    
    if MOCK_MODE:
        await asyncio.sleep(1) # Simulate inference latency
        demo_steps = [
            "Open a new browser tab.",
            "Create a blank document.",
            "Write the first sentence.",
            "Save the file.",
        ]
        if previous_failures == 1:
            return "Move your mouse to the browser icon."
        if previous_failures == 2:
            return "Put your hand on the mouse."
            
        total_steps = len(history) + len(skipped_history)
        return demo_steps[min(total_steps, len(demo_steps) - 1)]
            
    else:
        # If a task was rejected as too hard, route to MiniCPM to break it down further
        if previous_failures > 0 and rejected_task:
            print(f"TASK WAS REJECTED ('{rejected_task}'). ROUTING TO MINICPM FOR SUB-STEP BREAKDOWN...", flush=True)
            messages = [
                {"role": "system", "content": "You are a cognitive pacemaker. When a task is too hard, break it down into a single, even simpler, tiny physical starting action under 8 words. Return ONLY the starting action. CRITICAL: Focus strictly on the physical movement. Do NOT suggest thinking, planning, or remembering."},
                {"role": "user", "content": f"The task '{rejected_task}' was too hard. Break it down into a single, even simpler physical starting action."}
            ]
            fallback_res = await run_minicpm_chat(
                messages=messages,
                max_tokens=64, 
                temperature=0.1
            )
            raw_fallback = fallback_res['choices'][0]['message']['content'].strip()
            print(f"MINICPM BREAKDOWN OUT: {raw_fallback}", flush=True)
            content, _ = clean_and_validate_task(raw_fallback, [], [])
            return content or "Focus on the screen."

        # REAL INFERENCE using the exact string template used during fine-tuning
        system_msg = "You are a cognitive pacemaker. Break down goals into extremely tiny, atomic physical actions under 8 words."
        
        prompt = f"<extra_id_0>System\n{system_msg}\n\n"
        prompt += f"<extra_id_1>User\nGoal: {goal}\nCompleted Tasks: {history_str}\nFailures: {previous_failures}\nOutput the NEXT step.\n"
        prompt += f"<extra_id_1>Assistant\n"
        
        response = await run_nemotron(
            prompt,
            max_tokens=64, 
            temperature=0.3,
            stop=["\n", "<extra_id_1>", "Completed Tasks", "Goal:", "User:", "Assistant:", "<extra_id_"],
            grammar=grammar # Enforcing grammar here!
        )
        raw_content = response['choices'][0]['text'].strip()
        print(f"NEMOTRON RAW OUT: {raw_content}", flush=True)
        
        content, is_valid = clean_and_validate_task(raw_content, history, skipped_history)
        
        if not is_valid:
            print(f"NEMOTRON OUTPUT INVALID OR REPETITION ('{raw_content}'). FALLING BACK TO MINICPM...", flush=True)
            last_task = history[-1] if history else ""
            last_skipped = skipped_history[-1] if skipped_history else ""
            
            user_content = f"Goal: {goal}\nCompleted Tasks: {history_str}\nFailures: {previous_failures}\n"
            if last_skipped:
                user_content += (
                    f"CRITICAL: The user explicitly skipped '{last_skipped}'. Do NOT output it again. "
                    f"Stay on the current sub-topic or context of the last step, but generate an ALTERNATIVE physical starting action."
                )
            elif last_task:
                user_content += f"CRITICAL: Do NOT output '{last_task}'. Stay on the same sub-topic or context, but output the strictly NEXT new physical step.\n"
            else:
                user_content += "Output the NEXT step.\n"
            
            messages = [
                {"role": "system", "content": "You are a cognitive pacemaker. Focus strictly on the physical goal and break it down into an extremely tiny, atomic physical action under 8 words. Stay contextually coherent with the user's completed tasks, do NOT jump to unrelated sub-tasks, and output ONLY the single physical action step."},
                {"role": "user", "content": user_content}
            ]
            fallback_res = await run_minicpm_chat(
                messages=messages,
                max_tokens=64, 
                temperature=0.1
            )
            raw_fallback = fallback_res['choices'][0]['message']['content'].strip()
            print(f"MINICPM FALLBACK OUT: {raw_fallback}", flush=True)
            content, _ = clean_and_validate_task(raw_fallback, [], [])
            if not content:
                content = raw_fallback.strip('"\' ,.-1234567890)')
        
        return content or "Focus on the screen."


async def apply_activation_style(task: str, style="direct") -> str:
    if MOCK_MODE:
        await asyncio.sleep(0.5)
        if style == "calm":
            return f"When you are ready, {task[0].lower()}{task[1:]}"
        if style == "encouraging":
            return f"You can do this: {task}"
        return f"{task} Now."
        
    else:
        # REAL INFERENCE
        if style == "direct":
            return f"{task.capitalize()}."
            
        if style == "encouraging":
            system_prompt = (
                "You are a cognitive pacemaker. Rewrite the given task into an encouraging, supportive tone.\n"
                "RULES:\n"
                "1. Keep it as an action/command the user MUST do now. Start or end with an encouraging phrase like 'You got this!', 'Let's do it!', or 'Go ahead and...'.\n"
                "2. Do NOT write in the past tense, do NOT congratulate the user, and do NOT treat the task as already completed.\n"
                "3. Limit your response to EXACTLY ONE short sentence (under 12 words) and output ONLY the final rewritten task."
            )
        elif style == "calm":
            system_prompt = (
                "You are a cognitive pacemaker. Rewrite the given task into a calm, gentle tone.\n"
                "RULES:\n"
                "1. Keep it as an action/command. Use gentle prefixes like 'When you are ready, ...' or 'Take your time and ...'.\n"
                "2. Do NOT write in the past tense and do NOT congratulate the user.\n"
                "3. Limit your response to EXACTLY ONE short sentence (under 12 words) and output ONLY the final rewritten task."
            )
        else:
            system_prompt = (
                f"You are a cognitive pacemaker. Rewrite the given task into a {style} tone.\n"
                "RULES:\n"
                "1. Keep it as an action/command the user MUST do now.\n"
                "2. Do NOT write in the past tense, do NOT congratulate the user, and do NOT treat the task as already completed.\n"
                "3. Limit your response to EXACTLY ONE short sentence (under 12 words) and output ONLY the final rewritten task."
            )

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Task: {task}"}
        ]
        response = await run_minicpm_chat(
            messages=messages,
            max_tokens=40,
            temperature=0.2
        )
        content = response['choices'][0]['message']['content'].strip()
        print(f"MINICPM RAW OUT: {content}", flush=True)
        cleaned_content, is_valid = clean_and_validate_task(content, [], [])
        return cleaned_content or content.strip('"\' ,.-1234567890)')

# --- GRADIO UI ---

custom_css = """
:root, .dark {
    --body-background-fill: #0a0a0a !important;
    --background-fill-primary: #0a0a0a !important;
    --background-fill-secondary: #121212 !important;
    --border-color-primary: #1f2937 !important;
    --border-color-secondary: #374151 !important;
    --text-color-primary: #ffffff !important;
    --text-color-secondary: #d1d5db !important;
    --input-background-fill: #121212 !important;
    --input-border-width: 1px !important;
    --input-border-color: #1f2937 !important;
}

body {
    background-color: #0a0a0a !important;
    color: #ffffff !important;
    font-family: 'Helvetica Neue', sans-serif !important;
}

.gradio-container {
    background-color: #0a0a0a !important;
    border: none !important;
}

footer { display: none !important; }
.fog { opacity: 0.3; filter: blur(2px); transition: all 0.5s ease; }
.fade-in { animation: fadeIn 0.5s ease-in forwards; }
@keyframes fadeIn { from { opacity: 0; transform: translateY(10px); } to { opacity: 1; transform: translateY(0); } }

/* Goal Input Styling */
#goal-input textarea {
    background: transparent !important;
    border: none !important;
    border-bottom: 2px solid #4b5563 !important;
    color: white !important;
    font-size: 1.5rem !important;
    text-align: center !important;
    box-shadow: none !important;
    border-radius: 0 !important;
}
#goal-input textarea:focus { border-bottom: 2px solid white !important; }

/* Buttons */
#btn-start { background-color: #16a34a !important; color: white !important; font-weight: bold; font-size: 1.25rem !important; border: none !important; }
#btn-start:hover { background-color: #15803d !important; }
#btn-done { background-color: #16a34a !important; color: white !important; font-weight: bold; font-size: 1.25rem !important; }
#btn-skip { background-color: #4b5563 !important; color: white !important; font-weight: bold; font-size: 1.25rem !important; }
#btn-hard { background-color: #1f2937 !important; color: #d1d5db !important; font-weight: bold; font-size: 1.25rem !important; }
#btn-trace { background: transparent !important; border: none !important; color: #4b5563 !important; text-transform: uppercase; letter-spacing: 0.1em; font-size: 0.75rem !important; }
#btn-trace:hover { color: #d1d5db !important; }

/* Big Task Text */
#current-task { text-align: center; font-size: 3rem; font-weight: 800; min-height: 150px; display: flex; align-items: center; justify-content: center; color: #ffffff !important; }
@media (min-width: 768px) { #current-task { font-size: 4.5rem; } }
"""

def generate_breadcrumbs_html(history: list) -> str:
    html = "<div class='fog text-sm space-y-2'>"
    for t in history:
        html += f"<div>✓ {t}</div>"
    html += "</div>"
    return html

def generate_task_html(task: str, animate=True) -> str:
    class_str = "fade-in" if animate else ""
    return f"<div id='current-task' class='{class_str}'>{task}</div>"

async def process_step(state, action: str, goal: str = None, style: str = None, last_task: str = None):
    if not state:
        state = {}
    # Initialization
    if action == "start":
        state["goal"] = goal
        state["style"] = style
        state["current_task"] = ""
        state["current_raw_task"] = ""
        state["too_hard_count"] = 0
        state["history"] = []
        state["raw_history"] = []
        state["skipped_history"] = []
        state["trace"] = [{"event": "start", "goal": state["goal"], "style": state["style"]}]
        
    elif action == "too_hard":
        state["too_hard_count"] += 1
        append_trace(state, {
            "event": "too_hard",
            "task": state["current_task"],
            "too_hard_count": state["too_hard_count"],
        })
        
    elif action == "skip":
        append_trace(state, {"event": "skip", "task": state["current_task"]})
        if "skipped_history" not in state:
            state["skipped_history"] = []
        if state.get("current_raw_task"):
            state["skipped_history"].append(state["current_raw_task"])
        # Note: Do not increment too_hard_count
        
    elif action == "done":
        state["too_hard_count"] = 0
        if last_task:
            raw_task_done = state.get("current_raw_task") or last_task
            append_history(state, last_task, raw_task_done)
            append_trace(state, {"event": "done", "task": last_task})
    
    # HARD CIRCUIT BREAKER
    if state["too_hard_count"] >= 3:
        breaker_task = "You are out of activation energy. Step away from the screen for 3 minutes. I will be here."
        state["current_task"] = breaker_task
        append_trace(state, {"event": "circuit_breaker", "task": breaker_task})
        state["too_hard_count"] = 0 # Reset after breaking
        return state, generate_task_html(breaker_task), generate_breadcrumbs_html(state["history"][-3:]), gr.Row(visible=False)

    # Standard Loop
    rejected_task = state["current_task"] if state["too_hard_count"] > 0 else None
    
    raw_task = await generate_atomic_task(
        state["goal"],
        state["too_hard_count"],
        history=state["raw_history"],
        skipped_history=state.get("skipped_history", []),
        rejected_task=rejected_task,
    )
    styled_task = await apply_activation_style(raw_task, state["style"])
    state["current_task"] = styled_task
    state["current_raw_task"] = raw_task
    
    append_trace(state, {
        "event": "generated_step",
        "raw_task": raw_task,
        "styled_task": styled_task,
        "style": state["style"],
        "too_hard_count": state["too_hard_count"],
    })
    
    return state, generate_task_html(styled_task), generate_breadcrumbs_html(state["history"][-3:]), gr.Row(visible=True)

with gr.Blocks(css=custom_css, js="() => { document.documentElement.classList.add('dark'); }", title="Step-Zero") as app:
    session_state = gr.State()
    
    with gr.Column(visible=True) as screen_start:
        gr.HTML("<h1 class='text-4xl font-bold mb-8 text-center mt-20'>What is paralyzing you?</h1>")
        goal_input = gr.Textbox(elem_id="goal-input", show_label=False, placeholder="e.g. Write my thesis...", lines=1)
        
        with gr.Row(elem_classes="justify-center mt-8"):
            style_radio = gr.Radio(
                choices=["direct", "calm", "encouraging"], 
                value="direct", 
                show_label=False,
                container=False
            )
            
        start_btn = gr.Button("Start", variant="primary", elem_classes="mt-8 mx-auto w-48", elem_id="btn-start")
        
    with gr.Column(visible=False) as screen_task:
        breadcrumbs_display = gr.HTML(elem_id="breadcrumbs-container")
        task_display = gr.HTML(elem_id="task-container")
        
        with gr.Row(elem_id="controls-container") as controls_row:
            btn_done = gr.Button("I DID THIS", elem_id="btn-done")
            btn_skip = gr.Button("SKIP", elem_id="btn-skip")
            btn_hard = gr.Button("TOO HARD", elem_id="btn-hard")
            
        with gr.Row(elem_classes="justify-center mt-4 gap-4"):
            btn_trace = gr.Button("Export Trace", elem_id="btn-trace")
            btn_push = gr.Button("Push to Hub", elem_id="btn-trace")
            btn_reset = gr.Button("Start Over", elem_id="btn-trace")
            
        trace_download = gr.File(visible=False)
        hub_status = gr.HTML(visible=False, elem_classes="text-center text-sm text-gray-400 mt-2")

    # Temporary Loading State Function
    def show_loading():
        return gr.Column(visible=False), gr.Column(visible=True), generate_task_html("<span class='text-gray-600 animate-pulse'>Calculating constraint...</span>", animate=False), gr.Row(visible=False)

    def show_loading_step():
        return generate_task_html("<span class='text-gray-600 animate-pulse'>Loading next step...</span>", animate=False), gr.Row(visible=False)
        
    # --- Event Handlers ---
    
    async def handle_start(state, goal, style):
        return await process_step(state, "start", goal=goal, style=style)
        
    async def handle_done(state):
        if not state: state = {}
        return await process_step(state, "done", last_task=state.get("current_task"))
        
    async def handle_skip(state):
        return await process_step(state, "skip")
        
    async def handle_too_hard(state):
        return await process_step(state, "too_hard")
    
    # Start Session
    start_btn.click(
        fn=show_loading,
        outputs=[screen_start, screen_task, task_display, controls_row]
    ).then(
        fn=handle_start,
        inputs=[session_state, goal_input, style_radio],
        outputs=[session_state, task_display, breadcrumbs_display, controls_row]
    )

    goal_input.submit(
        fn=show_loading,
        outputs=[screen_start, screen_task, task_display, controls_row]
    ).then(
        fn=handle_start,
        inputs=[session_state, goal_input, style_radio],
        outputs=[session_state, task_display, breadcrumbs_display, controls_row]
    )

    # Mark Done
    btn_done.click(
        fn=show_loading_step,
        outputs=[task_display, controls_row]
    ).then(
        fn=handle_done,
        inputs=[session_state],
        outputs=[session_state, task_display, breadcrumbs_display, controls_row]
    )

    # Mark Skip
    btn_skip.click(
        fn=show_loading_step,
        outputs=[task_display, controls_row]
    ).then(
        fn=handle_skip,
        inputs=[session_state],
        outputs=[session_state, task_display, breadcrumbs_display, controls_row]
    )

    # Mark Too Hard
    btn_hard.click(
        fn=show_loading_step,
        outputs=[task_display, controls_row]
    ).then(
        fn=handle_too_hard,
        inputs=[session_state],
        outputs=[session_state, task_display, breadcrumbs_display, controls_row]
    )

    # Trace Export
    def export_trace_fn(state):
        import tempfile
        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
            json.dump(state.get("trace", []), f, indent=2)
            filepath = f.name
        return gr.File(value=filepath, visible=True)
        
    btn_trace.click(
        fn=export_trace_fn,
        inputs=[session_state],
        outputs=[trace_download]
    )

    # Push to Hub
    def push_to_hub_fn(state):
        import tempfile
        import time
        from huggingface_hub import HfApi
        
        try:
            if not state: state = {}
            with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
                json.dump(state.get("trace", []), f, indent=2)
                filepath = f.name
                
            api = HfApi()
            # Push to the user's dataset repo (tc043/step-zero-dataset as seen in README)
            # Using timestamp to avoid overwriting traces
            timestamp = int(time.time())
            api.upload_file(
                path_or_fileobj=filepath,
                path_in_repo=f"traces/trace_{timestamp}.json",
                repo_id="tc043/step-zero-dataset",
                repo_type="dataset"
            )
            return gr.HTML(value="<span class='text-green-500'>✓ Successfully pushed trace to HF Hub!</span>", visible=True)
        except Exception as e:
            return gr.HTML(value=f"<span class='text-red-500'>Error pushing to Hub: {str(e)}</span>", visible=True)

    btn_push.click(
        fn=push_to_hub_fn,
        inputs=[session_state],
        outputs=[hub_status]
    )

    def handle_reset():
        return {}, gr.Column(visible=True), gr.Column(visible=False), ""

    btn_reset.click(
        fn=handle_reset,
        outputs=[session_state, screen_start, screen_task, goal_input]
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)