import gradio as gr

from executor import KernelRequest, LocalCLIExecutor

SKILL_MD_DEFAULT = "skill_example/SKILL.md"

executor = LocalCLIExecutor()

theme = gr.themes.Default()

css = """
.orange-btn {
    background: #ea580c !important;
    color: white !important;
    border: none !important;
}
.orange-btn:hover {
    background: #c2410c !important;
}
.sidebar {
    border-right: 1px solid #e5e7eb !important;
    min-width: 200px !important;
    align-self: flex-start !important;
    position: sticky !important;
    top: 0 !important;
}
"""

# Holds the files dict from the latest run so callbacks can read them
_latest_files: dict[str, str] = {}


def generate_kernel(skills_path, prompt):
    global _latest_files
    _latest_files = {}
    request = KernelRequest(skill_path=skills_path, prompt=prompt)
    for response in executor.stream(request):
        if response.files:
            _latest_files = response.files
        choices = list(_latest_files.keys()) if _latest_files else []
        error_prefix = f"Error: {response.error}\n\n" if response.error else ""
        first_content = _latest_files.get(choices[0], "") if choices else ""
        yield (
            error_prefix + response.streaming_output,
            response.final_output,
            first_content,
            response.log,
            "Benchmark skipped",
            gr.update(choices=choices, value=choices[0] if choices else None),
        )


def on_file_select(filename):
    if filename and filename in _latest_files:
        return _latest_files[filename], gr.Tabs(selected="files")
    return "", gr.Tabs(selected="files")


with gr.Blocks(css=css) as demo:
    gr.Markdown(
        "# ⚡ Kernel Skills Sandbox\n"
        "Generate GPU kernels from natural language prompts. "
        "Point to your SKILL.md, describe what you need, and hit **Run**."
    )

    with gr.Group():
        skills_path = gr.Textbox(
            label="Path to SKILL.md",
            value=SKILL_MD_DEFAULT,
            placeholder="e.g. SKILL.md or /path/to/SKILL.md",
        )
        prompt = gr.Textbox(
            label="Prompt for Kernel Generation",
            lines=3,
            placeholder="Describe the kernel you want to generate...",
        )
        run_btn = gr.Button("Run", elem_classes=["orange-btn"])

    gr.Markdown("---")

    with gr.Row(equal_height=False):
        # LHS sidebar: file list
        with gr.Column(scale=1, min_width=200, elem_classes=["sidebar"]):
            gr.Markdown("**Generated Files**")
            file_dropdown = gr.Radio(
                label="Files",
                choices=[],
                interactive=True,
            )

        # RHS: tabbed content
        with gr.Column(scale=4):
            with gr.Tabs() as tabs:
                with gr.Tab("Output", id="output"):
                    streaming_output = gr.Textbox(
                        lines=20,
                        show_label=False,
                        interactive=False,
                    )
                with gr.Tab("Response", id="response"):
                    final_response = gr.Textbox(
                        lines=20,
                        show_label=False,
                        interactive=False,
                    )
                with gr.Tab("File Viewer", id="files"):
                    file_content = gr.Code(
                        label="File Content",
                        language=None,
                        lines=20,
                    )
                with gr.Tab("Live Log", id="log"):
                    live_log = gr.Textbox(
                        lines=20,
                        show_label=False,
                        interactive=False,
                    )
                with gr.Tab("Benchmark", id="benchmark"):
                    benchmark_results = gr.Textbox(
                        lines=20,
                        show_label=False,
                        interactive=False,
                    )

    run_btn.click(
        fn=generate_kernel,
        inputs=[skills_path, prompt],
        outputs=[streaming_output, final_response, file_content, live_log, benchmark_results, file_dropdown],
    )

    file_dropdown.change(
        fn=on_file_select,
        inputs=[file_dropdown],
        outputs=[file_content, tabs],
    )

demo.launch(theme=theme)