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 architecture diagram for {description} with the following specifications:
1. Canvas Setup:
- Use viewBox="0 0 800 600"

3. Visual Style Requirements:
- Use rounded rectangles (rx="10") for components
- Component size: 140px width, 80-120px height
- Text: centered, with title (14px, bold) and details (12px)
- Color scheme: professional color choices
- Consistent stroke width: 2px

4. Component Layout:
- Arrange components in [specify arrangement: grid/layered/hierarchical]
- Maintain clear spacing between components
- 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()
    # Wrap the generated code with the complete SVG tag and apply the dynamic viewBox
    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 full_svg_code, full_svg_code, "diagram.svg"

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, modified_svg_code, "diagram.svg"

with gr.Blocks() as demo:
    gr.Markdown("# Miru - Diagram Viewer")
    gr.Markdown("Generate a diagram by describing it, then edit the SVG code if needed and click **Update Diagram** to refresh the display and download.")

    with gr.Row():
        description_input = gr.Textbox(label="Describe your diagram or animation", placeholder="Enter a description...")
        generate_button = gr.Button("Generate Diagram")
    
    with gr.Row():
        svg_output = gr.HTML(label="Generated SVG", elem_id="svg-container")
        download_svg_output = gr.File(label="Download SVG")
    
    svg_code_output = gr.Textbox(label="SVG Code (editable)", lines=10, interactive=True)

    update_button = gr.Button("Update Diagram")

    generate_button.click(
        fn=generate_svg, 
        inputs=description_input, 
        outputs=[svg_output, svg_code_output, download_svg_output]
    )
    
    update_button.click(
        fn=update_svg, 
        inputs=svg_code_output, 
        outputs=[svg_output, svg_code_output, download_svg_output]
    )

demo.launch()