import gradio as gr
from jinja2 import Template
import json

def render_mermaid(jinja_text, variables):
    try:
        template = Template(jinja_text)
        rendered_code = template.render(**variables)
        mermaid_md = f"```mermaid\n{rendered_code}\n```"
        return mermaid_md
    except Exception as e:
        return f"**Template Error:** {str(e)}"

def add_variable(key, value, current_vars):
    if not key:
        return current_vars, "**Variable name is required**"
    if key in current_vars:
        return current_vars, f"⚠️ Variable '{key}' already exists. Use update instead."
    try:
        parsed_value = json.loads(value)
    except json.JSONDecodeError:
        parsed_value = value
    current_vars[key] = parsed_value
    return current_vars, f"✅ Variable '{key}' added"

def update_variable(key, value, current_vars):
    if not key:
        return current_vars, "**Variable name is required**"
    if key not in current_vars:
        return current_vars, f"⚠️ Variable '{key}' does not exist. Use add instead."
    try:
        parsed_value = json.loads(value)
    except json.JSONDecodeError:
        parsed_value = value
    current_vars[key] = parsed_value
    return current_vars, f"♻️ Variable '{key}' updated"

def delete_variable(key, current_vars):
    if key in current_vars:
        del current_vars[key]
        return current_vars, f"🗑️ Variable '{key}' deleted"
    return current_vars, f"⚠️ Variable '{key}' not found"

def reset_variables():
    return {}

default_template = """
flowchart TD
    Start(["Start"]) --> A{"Is city known?"}
    A -->|Yes| B["You are in {{ user_city }}"]
    A -->|No| C["Please choose a city"]
"""

custom_css = """
.gr-box { padding: 1.5rem; }

#jinja-code-box {
    max-height: 400px;
    overflow-y: auto;
    border: 1px solid #ddd;
    border-radius: 6px;
}

.markdown-preview {
    padding: 1rem;
    max-height: 85vh;
    overflow-y: auto;
    white-space: pre-wrap;
    font-size: 0.95rem;
    line-height: 1.5;
}
.gr-button { margin-right: 0.5rem; }
"""

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("## 🧠 Jinja2 → Mermaid Renderer (Dynamic Variables - Interactive)")

    with gr.Row():
        with gr.Column(scale=7):
            jinja_input = gr.Code(
                label="✍️ Jinja2 + Mermaid Template",
                language="html",
                value=default_template,
                lines=18,
                elem_id="jinja-code-box"
            )

            gr.Markdown("### ➕ Add / ♻️ Update / 🗑️ Delete Template Variables")
            with gr.Row():
                with gr.Column(scale=2):
                    var_key = gr.Textbox(label="🆔 Variable Name", placeholder="e.g., user_city")
                    var_value = gr.Textbox(label="📦 Value (JSON)", placeholder='e.g., "London"')
                    
                    with gr.Row():
                        add_btn = gr.Button("➕ Add Variable")
                        update_btn = gr.Button("♻️ Update Variable")
                    with gr.Row():
                        delete_btn = gr.Button("🗑️ Delete Variable")
                        reset_btn = gr.Button("🔁 Reset Variables")

                with gr.Column(scale=2):
                    error_box = gr.Markdown("", visible=False)
                    rendered_vars = gr.JSON(label="📋 Current Variables")
                    variables_state = gr.State({})

        with gr.Column(scale=5):
            gr.Markdown("### 📈 Live Mermaid Preview")
            rendered_md = gr.Markdown(elem_classes="markdown-preview")

    add_btn.click(
        fn=add_variable,
        inputs=[var_key, var_value, variables_state],
        outputs=[variables_state, error_box],
        show_progress=False
    ).then(
        fn=lambda x: x,
        inputs=variables_state,
        outputs=rendered_vars
    ).then(
        fn=render_mermaid,
        inputs=[jinja_input, variables_state],
        outputs=rendered_md
    )

    update_btn.click(
        fn=update_variable,
        inputs=[var_key, var_value, variables_state],
        outputs=[variables_state, error_box],
        show_progress=False
    ).then(
        fn=lambda x: x,
        inputs=variables_state,
        outputs=rendered_vars
    ).then(
        fn=render_mermaid,
        inputs=[jinja_input, variables_state],
        outputs=rendered_md
    )

    delete_btn.click(
        fn=delete_variable,
        inputs=[var_key, variables_state],
        outputs=[variables_state, error_box],
        show_progress=False
    ).then(
        fn=lambda x: x,
        inputs=variables_state,
        outputs=rendered_vars
    ).then(
        fn=render_mermaid,
        inputs=[jinja_input, variables_state],
        outputs=rendered_md
    )

    reset_btn.click(
        fn=reset_variables,
        inputs=[],
        outputs=variables_state
    ).then(
        fn=lambda x: x,
        inputs=variables_state,
        outputs=rendered_vars
    ).then(
        fn=render_mermaid,
        inputs=[jinja_input, variables_state],
        outputs=rendered_md
    )

    jinja_input.change(render_mermaid, inputs=[jinja_input, variables_state], outputs=rendered_md)

    demo.load(render_mermaid, inputs=[jinja_input, variables_state], outputs=rendered_md)

demo.launch()