import spaces  # MUST be first — before any torch/CUDA import
import torch
import json
import re
import gradio as gr

from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MODEL_ID = "caid-technologies/parti-base"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to("cuda")
model.eval()

SYSTEM_PROMPT = (
    "You design hobbyist electronics projects. Given a request, reply with a single "
    "JSON object describing the full project. Output only the JSON."
)

MODE_PROMPTS = {
    "Full project plan": "",
    "Parts list only": " Respond with ONLY the parts list (a JSON array of components).",
    "Wiring map only": " Respond with ONLY the wiring/connection map (a JSON object of connections).",
    "Build steps only": " Respond with ONLY the ordered build steps (a JSON array of steps).",
}


def _build_messages(user_prompt: str, mode: str) -> list:
    suffix = MODE_PROMPTS.get(mode, "")
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_prompt + suffix},
    ]


def _try_format_json(text: str) -> str:
    """Attempt to extract and pretty-print a JSON block from the model output."""
    text = text.strip()
    # Try direct parse
    try:
        parsed = json.loads(text)
        return json.dumps(parsed, indent=2)
    except (json.JSONDecodeError, ValueError):
        pass
    # Try to find a JSON object or array in the text
    for pattern in [r"\{.*\}", r"\[.*\]"]:
        match = re.search(pattern, text, re.DOTALL)
        if match:
            try:
                parsed = json.loads(match.group())
                return json.dumps(parsed, indent=2)
            except (json.JSONDecodeError, ValueError):
                continue
    return text


@spaces.GPU(duration=75)
def generate(
    project_idea: str,
    mode: str,
    max_tokens: int,
    repetition_penalty: float,
):
    """Generate a structured hardware project plan from a plain-English description.

    Args:
        project_idea: A description of the hardware project you want to build.
        mode: What to generate — full plan, parts list, wiring map, or build steps.
        max_tokens: Maximum number of new tokens to generate.
        repetition_penalty: Penalty for repeated tokens (1.1 recommended for wiring lists).
    """
    messages = _build_messages(project_idea, mode)
    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to("cuda")

    text_streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    generation_kwargs = dict(
        **inputs,
        max_new_tokens=max_tokens,
        do_sample=False,
        repetition_penalty=repetition_penalty,
        pad_token_id=tokenizer.eos_token_id,
        streamer=text_streamer,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    collected = ""
    for text in text_streamer:
        collected += text
        yield collected

    thread.join()

    # Try to pretty-print the JSON if valid
    yield _try_format_json(collected)


CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Hardware Blueprint Generator") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Hardware Blueprint Generator\n"
            "Describe a hardware project and get a structured JSON blueprint — "
            "parts list, wiring map, build steps, and design check.\n\n"
            "Powered by [caid-technologies/blueprint-base](https://huggingface.co/caid-technologies/blueprint-base), "
            "a Qwen2.5-3B fine-tune for hobbyist electronics project planning."
        )

        with gr.Row():
            project_input = gr.Textbox(
                label="Project Description",
                placeholder="e.g. A compact desk clock with an e-ink display and an IR remote.",
                lines=3,
                scale=4,
            )
            mode_selector = gr.Dropdown(
                label="Output Mode",
                choices=list(MODE_PROMPTS.keys()),
                value="Full project plan",
                scale=1,
            )

        generate_btn = gr.Button("Generate Blueprint", variant="primary")

        output = gr.Textbox(
            label="Generated Blueprint (JSON)",
            lines=25,
        )

        with gr.Accordion("Advanced Settings", open=False):
            max_tokens = gr.Slider(
                label="Max New Tokens",
                minimum=1024,
                maximum=8192,
                value=6144,
                step=512,
                info="Higher values allow longer plans but take more time.",
            )
            rep_penalty = gr.Slider(
                label="Repetition Penalty",
                minimum=1.0,
                maximum=2.0,
                value=1.1,
                step=0.05,
                info="1.1 recommended — prevents wiring lists from repeating.",
            )

        gr.Examples(
            examples=[
                ["A compact desk clock with an e-ink display and an IR remote.", "Full project plan"],
                ["A weather station with temperature, humidity, and pressure sensors that uploads data to a web dashboard.", "Full project plan"],
                ["A pet feeder that dispenses food on a schedule and sends phone notifications.", "Parts list only"],
                ["A LoRa-based air quality monitor with a solar panel for power.", "Build steps only"],
            ],
            inputs=[project_input, mode_selector],
            outputs=output,
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.Markdown(
            "---\n"
            "**Note:** This is an early research preview. Output is a *first draft* to review — "
            "not a finished engineering design. Always sanity-check before building."
        )

    generate_btn.click(
        fn=generate,
        inputs=[project_input, mode_selector, max_tokens, rep_penalty],
        outputs=output,
        api_name="generate",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)