import os
import gradio as gr
import requests
from groq import Groq
import cairosvg

groq_api_key = os.getenv("MiruL")

def generate_svg(description):
    client = Groq(api_key=groq_api_key)
    completion = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[
            {"role": "system", "content": """
Create an SVG diagram for {description} with the following specifications:
1. Canvas Setup:
- Use viewBox="0 0 800 600"
2. Visual Style Requirements (apply when needed):
- Text: centered, with title (14px, bold) and details (12px)
- Ensure consistent spacing
- Ensure proper use of arrow only when needed
- Color scheme: professional color choices
- Consistent stroke width: 2px
- Avoid line crossings in connections
3. Additional Requirements:
Put the name of the Diagram in bold on top
- Multiple Color coding
Please generate a complete SVG code that can be directly used. Ensure the code follows SVG best practices and includes proper spacing and formatting for readability. 
Make sure it's just the code only
Generate the correct code, ensure it is free from construct error and this kind of error: Unescaped '<' not allowed in attributes values
"""}, 
            {"role": "user", "content": description}
        ],
        temperature=0.3,
        max_completion_tokens=4096,
        top_p=1,
        stream=False,
        stop=None,
    )
    
    svg_code = completion.choices[0].message.content.strip()
    full_svg_code = f"<svg xmlns='http://www.w3.org/2000/svg' width='100%' height='100%' viewBox='0 0 1000 400'>{svg_code}</svg>"

        # Save the SVG code to a file for record and future editing
    with open("diagram.svg", "w") as f:
        f.write(full_svg_code)
        
    # Convert the SVG code to PNG using CairoSVG
    png_filename = "diagram.png"
    cairosvg.svg2png(bytestring=full_svg_code.encode('utf-8'), write_to=png_filename)
    
    # Return the PNG for display, the PNG file for download, and the SVG code for editing later.
    return full_svg_code, png_filename, full_svg_code

def update_svg(modified_svg_code):
    # Save the modified SVG code to the file
    with open("diagram.svg", "w") as f:
        f.write(modified_svg_code)
    png_filename = "diagram.png"
    cairosvg.svg2png(bytestring=modified_svg_code.encode('utf-8'), write_to=png_filename)
    return png_filename, png_filename

with gr.Blocks(css="""

""") as demo:
    gr.Markdown("# Miru - Diagram Generator and Viewer")
    gr.Markdown("Use the **Generate Diagram** tab to create a diagram by describing it. In the **Update Diagram** tab, you can edit the generated SVG code and update the diagram.")

    # State to store the generated SVG code for editing in Tab 2
    svg_state = gr.State()

    with gr.Tabs():
        with gr.TabItem("Generate Diagram"):
            description_input = gr.Textbox(label="Describe your diagram or animation",
                                           placeholder="Enter a description...",
                                           interactive=True)
            # Generated SVG displayed in a larger HTML area
            svg_output = gr.HTML(label="Generated SVG", elem_id="svg-container", value="")
            download_svg_output = gr.File(label="Download SVG")
            
            # Submit on pressing Enter in the description box
            description_input.submit(fn=generate_svg,
                                     inputs=description_input,
                                     outputs=[svg_output, download_svg_output, svg_state])
        
        with gr.TabItem("Update Diagram"):
            gr.Markdown("Edit the SVG code below and click **Update Diagram** to refresh the display and download file.")
            modified_svg_input = gr.Textbox(label="Edit SVG Code", lines=10, interactive=True)
            load_button = gr.Button("Load Generated Code")
            update_button = gr.Button("Update Diagram")
            updated_svg_output = gr.HTML(label="Updated SVG", elem_id="updated-svg-container", value="")
            updated_download_svg_output = gr.File(label="Download SVG")
            
            # Load the generated SVG code from state into the editable textbox
            load_button.click(lambda code: code, inputs=svg_state, outputs=modified_svg_input)
            # Update the diagram when the update button is clicked
            update_button.click(fn=update_svg,
                                inputs=modified_svg_input,
                                outputs=[updated_svg_output, updated_download_svg_output])

demo.launch()