import os
from pathlib import Path

os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")

import gradio as gr

from wisp_inference import MODEL_ID, MODEL_REVISION, MODES, complete
from visuals import make_theme
from garden_design import CSS, FOOTER, INTRO, INTRO_JS, PLAYGROUND, RESULTS

gr.set_static_paths(paths=[Path(__file__).parent / "assets"])


def generate(prefix, suffix, mode, max_new_tokens):
    try:
        return complete(prefix, suffix, mode, max_new_tokens)
    except ValueError as error:
        raise gr.Error(str(error)) from error


with gr.Blocks(title="Wisp Coder — Code at the Cursor") as demo:
    gr.HTML(INTRO, elem_id="intro", apply_default_css=False, js_on_load=INTRO_JS)
    with gr.Column(elem_id="workbench", min_width=0):
        gr.HTML(PLAYGROUND, elem_id="playground-heading", apply_default_css=False)
        with gr.Row(elem_id="options"):
            mode = gr.Radio(MODES, value=MODES[0], label="Completion mode", scale=3)
            limit = gr.Slider(16, 128, value=64, step=16, label="Maximum new tokens", scale=2)
        with gr.Row(elem_id="editors"):
            prefix = gr.Textbox(
                label="Before the cursor", lines=7, max_lines=12, elem_classes=["code-box"],
                value="def is_even(n):\n    return ", interactive=True,
            )
            suffix = gr.Textbox(
                label="After the cursor · fill-in-the-middle only", elem_classes=["code-box"],
                lines=7, max_lines=12, value="\n", interactive=True,
                placeholder="Optional code that comes after your suggestion…",
            )
        gr.Examples(
            examples=[
                ["def is_even(n):\n    return ", "\n", MODES[0], 32],
                ["def square(x):\n    return ", "\n", MODES[0], 32],
                ["def factorial(n):\n", "", MODES[1], 64],
                ["def is_even(n):\n    ", "", MODES[1], 32],
            ],
            inputs=[prefix, suffix, mode, limit],
            label="Try a starting point — then generate",
            example_labels=["Even or odd · fill", "Square · fill", "Factorial · continue", "Even or odd · continue"],
            elem_id="seed-examples",
            cache_examples=False,
        )
        with gr.Row(elem_id="run-row"):
            run = gr.Button("Generate suggestion  ↗", variant="primary", elem_id="run", scale=0)
            status = gr.Markdown("Ready to explore. Pick a starting point or edit the code above.", elem_id="generation-status")
    with gr.Column(elem_id="harvest", min_width=0):
        gr.HTML(RESULTS, elem_id="results-heading", apply_default_css=False)
        with gr.Row(elem_id="outputs"):
            suggestion = gr.Textbox(label="Generated suggestion", lines=8, max_lines=16,
                                    placeholder="Your model's suggestion will appear here.",
                                    elem_classes=["code-box"], interactive=False, buttons=["copy"])
            combined = gr.Textbox(label="Code with suggestion inserted", lines=8, max_lines=16,
                                  placeholder="The full snippet, with the suggestion in place.",
                                  elem_classes=["code-box"], interactive=False, buttons=["copy"])
    run.click(generate, inputs=[prefix, suffix, mode, limit],
              outputs=[suggestion, combined, status], api_name="complete",
              concurrency_limit=1)
    with gr.Accordion("About this research demo", open=False, elem_id="about"):
        gr.Markdown(
            "Wisp was trained with fill-in-the-middle and a native multi-token-prediction module. "
            "This demo runs the standard Llama trunk using Transformers on CPU; it does not load "
            "the MTP sidecar. The full release has 108.2M parameters including that module.\n\n"
            "This is a base code-completion model. It has limited completion quality and does not "
            "follow chat instructions reliably. Suggestions can be incomplete, incorrect, or insecure; "
            "review them before use. Generated code is displayed only and is never executed here.\n\n"
            "The registered short-hole evaluation found no superiority over the compared FIM models. "
            f"[Read the results](https://huggingface.co/{MODEL_ID}/blob/main/fim-competitive-v1/FIM_RESULTS.md).\n\n"
            "For example, a local greedy smoke generated `1 + x**2` for `square(x)`, "
            "which is incorrect. Example prompts are demonstrations, not a quality benchmark.\n\n"
            "This demo caps input at 1,536 tokens and output at 128 tokens and uses deterministic "
            "greedy decoding. In Continue code mode, the suffix is ignored. Inputs are sent to this "
            "Hugging Face Space for inference. The app does not save prompts or outputs.\n\n"
            f"Model revision: `{MODEL_REVISION}`. The displayed generation time excludes model loading "
            "and queue time and is not an Apple Silicon benchmark."
        )
    gr.HTML(FOOTER, elem_id="wisp-footer", apply_default_css=False)

demo.queue(max_size=8, default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch(theme=make_theme(), css=CSS)