import gradio as gr

def beautify_diagram_types(diagram_types):
    first_diagram = "use-case-diagram"
    if first_diagram in diagram_types:
        diagram_types.remove(first_diagram)
    diagram_types.insert(0, first_diagram)
    diagram_types.sort()
    return [name.replace("-", " ").title() for name in diagram_types]

def create_interface(diagram_types, respond):
    diagram_types = beautify_diagram_types(diagram_types)
    with gr.Blocks(title="PlantUML Agent", theme=gr.themes.Soft()) as demo:
        gr.HTML(
            """
            <h1 style="text-align: center; color: white;">PlantUML Agent</h1>
            <div style="text-align: center; margin-bottom: 20px;">
                <img src="https://repository-images.githubusercontent.com/553868400/2f7375d8-c2ee-44fe-81d0-a8a29ec284da" alt="PlantUML Logo" style="height: 60px;">
            </div>
            <p style="text-align: center; font-size: 1.1em; color: white;">
                Generate UML diagrams using <a href="https://plantuml.com" target="_blank" style="color: #1A73E8; text-decoration: none;">PlantUML</a> syntax with the help of an AI agent.
            </p>
            <hr style="border: 0; height: 1px; background: #eee; margin: 20px 0;">
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                
                gr.Markdown("## 🔑 Mistral API Key Input")
                api_key_input = gr.Textbox(
                    label="🔐 Mistral API Key",
                    placeholder="Input your Mistral API key here",
                    type="password",
                    lines=1,
                    interactive=True
                )

                gr.Markdown("## 📂 Diagram Type")
                diagram_dropdown = gr.Dropdown(
                    choices=diagram_types,
                    value="Use Case Diagram",
                    label="Choose the type of UML diagram",
                    interactive=True
                )
                
                gr.Markdown("## 📝 Input Description")
                msg = gr.Textbox(
                    label="Describe your UML diagram",
                    placeholder="e.g., create a basic use case class",
                    value="Create a basic use case diagram",
                    lines=3,
                    interactive=True
                )
                
                gr.Markdown("## 🖼️ Image (optional)")
                image_input = gr.Image(
                    label="Upload an image to assist in diagram generation",
                    type="filepath",
                    interactive=True,
                    height=200,
                )

                submit = gr.Button("🚀 Generate Diagram", variant="primary", size="lg")

            with gr.Column(scale=2):
                gr.Markdown("## 📊 Generated Diagram")

                gr.Markdown("### PlantUML Code")
                diagram_text_output = gr.Code(
                    label="The generated PlantUML code",
                    lines=15,
                    interactive=False,
                )

                gr.Markdown("### Diagram Preview")
                diagram_output = gr.HTML(
                    label="Visual representation of your diagram",
                    elem_id="diagram-preview"
                )

        gr.HTML(
            """
            <hr style="border: 0; height: 1px; background: #eee; margin: 30px 0 15px 0;">
            <div style='text-align: center; font-size: 0.9em; color: #777;'>
            Powered by Mistral, BM25Retriever, LangChain, Docling, LRUCache, Gradio, and PlantUML.
            </div>
            <div style='text-align: center; font-size: 0.9em; color: #777; margin-top: 5px;'>
            Created by <a href='https://linkedin.com/in/savi8sant8s' target='_blank' style="color: #1A73E8; text-decoration: none;">Sávio Santos</a>
            </div>
            """
        )

        submit.click(respond, [msg, diagram_dropdown, image_input, api_key_input], [msg, diagram_text_output, diagram_output])
        msg.submit(respond, [msg, diagram_dropdown, image_input, api_key_input], [msg, diagram_text_output, diagram_output])

    return demo