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

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"
3. Visual Style Requirements (apply when needed):
- Text: centered, with title (14px, bold) and details (12px)
- Color scheme: professional color choices
- Consistent stroke width: 2px
- Avoid line crossings in connections
6. 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.7,
        max_completion_tokens=1024,
        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 download
    with open("diagram.svg", "w") as f:
        f.write(full_svg_code)
    
    # Return the HTML display, download file, and store the full code in state for editing in Tab 2.
    return full_svg_code, "diagram.svg", 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)
    return modified_svg_code, "diagram.svg"

with gr.Blocks() as demo:
    gr.Markdown("# Miru - Diagram 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="")
            svg_output.style(height="600px")
            download_svg_output = gr.File(label="Download SVG")
            
            # When the user presses Enter, generate_svg is called and updates the outputs and state.
            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="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()