main.py

import gradio as gr
from agent import session_service, root_agent, APP_NAME, runner, COMPONENT_INFO_DICT
from google.genai import types
from pprint import pprint, pformat
from uuid import uuid4
import os
import tempfile


def generate_mermaid(components, wires=[], direction: str = "LR") -> str:
    """
    Build a Mermaid-JS flow-chart for a simple “high-level” schematic.

    Parameters
    ----------
    components : list[dict]
        Each dict needs the keys
            • 'component_name' – human-readable label (e.g. "Arduino Uno R3")
            • 'id'             – unique node ID used in the diagram
    wires : list[dict]
        Each dict needs the keys
            • 'start_id', 'start_pin' – origin node & pin
            • 'end_id',   'end_pin'   – destination node & pin
    direction : str, optional
        Flow direction for the chart header; one of "LR", "RL", "TB", "TD".
        Defaults to "LR" (left-to-right).

    Returns
    -------
    str
        A Markdown-fenced Mermaid block you can drop straight into GitHub,
        Obsidian, Notion, etc.
    """

    mermaid = [f"```mermaid", f"flowchart {direction}"]

    # 1) Declare nodes
    for node_id, node_name in components.items():
        label = f"{node_name}\n({node_id})"
        mermaid.append(f'    {node_id}["{label}"]')

    # 2) Draw labelled connections
    for wire in wires:
        start = wire["start_id"]
        end = wire["end_id"]
        label = f'{wire["start_pin"]} → {wire["end_pin"]}'
        mermaid.append(f'    {start} -- "{label}" --> {end}')

    mermaid.append("```")
    return "\n".join(mermaid)


async def interact_with_agent(
    prompt,
    history: list[gr.ChatMessage],
    session_id: gr.State,
    local_storage: gr.BrowserState,
    request: gr.Request,
):
    messages = []
    messages.append(
        gr.ChatMessage(
            role="assistant", content="Thinking...", metadata={"status": "pending"}
        )
    )
    yield messages, None, None, None
    # yield messages

    user_id = local_storage[0]

    print(user_id)
    print(session_id)

    session = session_service.get_session(
        app_name=APP_NAME, user_id=user_id, session_id=session_id
    )

    if session is None:
        session = session_service.create_session(
            app_name=APP_NAME, user_id=user_id, session_id=session_id
        )

    fetch_code = lambda: session_service.get_session(
        app_name=APP_NAME, user_id=user_id, session_id=session_id
    ).state.get("code")
    fetch_components = lambda: session_service.get_session(
        app_name=APP_NAME, user_id=user_id, session_id=session_id
    ).state.get("components")
    fetch_wires = lambda: session_service.get_session(
        app_name=APP_NAME, user_id=user_id, session_id=session_id
    ).state.get("wires")

    def mermaid_assemble():
        components = fetch_components()
        wires = fetch_wires()
        if components == None:
            return gr.Markdown("No components yet")
        if wires == None:
            wires = []
        mermaid = generate_mermaid(components, wires)
        return gr.Markdown(mermaid)

    content = types.Content(role="user", parts=[types.Part(text=prompt)])

    temp = None

    async for event in runner.run_async(
        user_id=user_id, session_id=session_id, new_message=content
    ):
        """
        print(
            f"  [Event] Author: {event.author}, Type: {type(event).__name__}, Final: {event.is_final_response()}, Content: {event.content}"
        )
        """
        for call in event.get_function_calls():
            messages.append(
                gr.ChatMessage(
                    role="assistant",
                    content=(
                        f"Invoking {call.name} with args: \n {pformat(call.args)}"
                        if not call.name == "write_code"
                        else gr.Code(call.args["code"], language="cpp")
                    ),
                    metadata={"title": f"🛠️ Used tool {call.name}"},
                )
            )
            yield messages, None, None, None
        for call in event.get_function_responses():
            if call.name == "compile_with_arduino" and call.response["success"]:
                path = call.response["project_dir"]
                elf_path = os.path.join(path, "output", "arduino_project.ino.elf")
                print(f"elf_path: {elf_path}")
                os.chmod(elf_path, 0o777)

                yield messages, None, None, gr.DownloadButton(
                    "Download ELF",
                    elf_path,
                    variant="primary",
                    render=False,
                    interactive=True,
                )
        if event.is_final_response():
            messages.append(
                gr.ChatMessage(role="assistant", content=event.content.parts[0].text)
            )
        yield messages, gr.Code(
            language="cpp", value=fetch_code()
        ), mermaid_assemble(), (
            None
            if not temp
            else gr.DownloadButton(
                "Download ELF",
                temp,
                variant="primary",
                render=False,
                interactive=True,
            )
        )

    session = session_service.get_session(
        app_name=APP_NAME, user_id=user_id, session_id=session_id
    )

    pprint(session.state)


with gr.Blocks() as demo:
    session_id = gr.State(str(uuid4()))
    code = gr.Code(language="cpp", render=False)
    markdown_schematic = gr.Markdown(render=False)
    local_storage = gr.BrowserState([str(uuid4())])
    download_button = gr.DownloadButton(
        "Download ELF",
        "/tmp/arduino_project/output/arduino_project.ino.elf",
        interactive=False,
        render=False,
    )
    COMPONENT_NAMES_MD = "\n".join([f'- {name}' for name in COMPONENT_INFO_DICT.keys()])

    with gr.Row():
        gr.Markdown(
            f"""
# Arduino Hardware and Software Vibecoder

This is a prototype of a vibecoder (like Lovable) that can generate Arduino code and schematics for you.
Given a prompt, it well look up the components it has availble and generate a high-level schematic and code.
It compiles the code to make sure it works and then gives you a download link for the ELF file that can then be uploaded to the Arduino.

It has the following components:
{COMPONENT_NAMES_MD}
                    """
        )
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot(label="Agent", type="messages")

            def clear_chat():
                return [
                    str(uuid4()),
                    gr.Code(language="cpp", render=False),
                    gr.Markdown(render=False),
                    gr.DownloadButton(
                        "Download ELF",
                        "/tmp/arduino_project/output/arduino_project.ino.elf",
                        render=False,
                        interactive=False,
                    ),
                ]

            chatbot.clear(
                lambda: clear_chat(),
                outputs=[session_id, code, markdown_schematic, download_button],
            )
            chat_interface = gr.ChatInterface(
                interact_with_agent,
                chatbot=chatbot,
                type="messages",
                examples=[
                    [
                        'Design a schematic and write code for an arduino driven "Joke of the day machine" where the user presses a physical button and then a random joke of the day is displayed on the display.'
                    ],
                    ["Write a Hello World program for an Arduino Uno"],
                    ["Write a program to blink an LED on an Arduino Uno"],
                ],
                additional_inputs=[session_id, local_storage],
                additional_outputs=[code, markdown_schematic, download_button],
                fill_height=True,
                show_progress="full",
            )
        with gr.Column():
            code.render()
            markdown_schematic.render()

    with gr.Row():
        download_button.render()

    # mm
demo.launch(show_api=False)