import gradio as gr
import vtracer
import os

def convert_to_vector(
    image, 
    colormode="color", 
    hierarchical="stacked", 
    mode="spline", 
    filter_speckle=4, 
    color_precision=6, 
    layer_difference=16, 
    corner_threshold=60, 
    length_threshold=4.0, 
    max_iterations=10, 
    splice_threshold=45, 
    path_precision=3
):
    input_path = "temp_input.jpg"
    output_path = "svg_output.svg"

    # Save the input image to a temporary file
    image.save(input_path)

    # Convert the image to SVG using VTracer
    vtracer.convert_image_to_svg_py(
        input_path,
        output_path,
        colormode=colormode,
        hierarchical=hierarchical,
        mode=mode,
        filter_speckle=int(filter_speckle),
        color_precision=int(color_precision),
        layer_difference=int(layer_difference),
        corner_threshold=int(corner_threshold),
        length_threshold=float(length_threshold),
        max_iterations=int(max_iterations),
        splice_threshold=int(splice_threshold),
        path_precision=int(path_precision)
    )

    # Read the SVG output
    with open(output_path, "r") as f:
        svg_content = f.read()

    # Return the SVG file path and content
    return gr.HTML(f'<svg viewBox="0 0 {image.width} {image.height}">{svg_content}</svg>'), output_path

def handle_color_mode(value):
    return value

def clear_inputs():
    return None, "color", "stacked", "spline", 4, 6, 16, 60, 4.0, 10, 45, 3

def update_interactivity_and_visibility(colormode):
    is_color_mode = colormode == "color"
    return (
        gr.update(interactive=is_color_mode),
        gr.update(interactive=is_color_mode),
        gr.update(visible=is_color_mode)
    )

css = """
#col-container {
    margin: 0 auto;
    max-width: 960px;
}
.generate-btn {
    background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%) !important;
    border: none !important;
    color: white !important;
}
.generate-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
}
"""

# Define the Gradio interface
with gr.Blocks(css=css) as app:
    with gr.Column(elem_id="col-container"):
        gr.HTML("""
            <div style="text-align: center;">
                <h2>Image to Vector Converter ⚡</h2>
                <p>Converts raster images (JPG, PNG, WEBP) to vector graphics (SVG).</p>
            </div>
        """)
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="Upload Image")
                convert_button = gr.Button("✨ Convert to SVG", variant='primary', elem_classes=["generate-btn"])
                clear_button = gr.Button("Clear Inputs", variant='secondary')

            with gr.Column():
                html = gr.HTML(label="SVG Output")
                svg_output = gr.File(label="Download SVG")

        with gr.Row():
            colormode = gr.Radio(choices=["color", "binary"], value="color", label="Color Mode")
            hierarchical = gr.Radio(choices=["stacked", "cutout"], value="stacked", label="Hierarchical Mode")
            mode = gr.Radio(choices=["spline", "polygon"], value="spline", label="Mode")

        with gr.Row():
            filter_speckle = gr.Slider(minimum=0, maximum=20, value=4, label="Filter Speckle")
            color_precision = gr.Slider(minimum=1, maximum=10, value=6, label="Color Precision")
            layer_difference = gr.Slider(minimum=1, maximum=32, value=16, label="Layer Difference")
            corner_threshold = gr.Slider(minimum=0, maximum=180, value=60, label="Corner Threshold")
            length_threshold = gr.Slider(minimum=0, maximum=10, value=4.0, label="Length Threshold")
            max_iterations = gr.Slider(minimum=1, maximum=20, value=10, label="Max Iterations")
            splice_threshold = gr.Slider(minimum=0, maximum=90, value=45, label="Splice Threshold")
            path_precision = gr.Slider(minimum=1, maximum=10, value=3, label="Path Precision")

    # Event handlers
    colormode.change(handle_color_mode, inputs=colormode, outputs=colormode)
    hierarchical.change(handle_color_mode, inputs=hierarchical, outputs=hierarchical)
    mode.change(handle_color_mode, inputs=mode, outputs=mode)

    colormode.change(
        update_interactivity_and_visibility,
        inputs=colormode,
        outputs=[color_precision, layer_difference, hierarchical]
    )

    clear_button.click(
        clear_inputs,
        outputs=[
            image_input,
            colormode,
            hierarchical,
            mode,
            filter_speckle,
            color_precision,
            layer_difference,
            corner_threshold,
            length_threshold,
            max_iterations,
            splice_threshold,
            path_precision
        ]
    )

    convert_button.click(
        convert_to_vector,
        inputs=[
            image_input,
            colormode,
            hierarchical,
            mode,
            filter_speckle,
            color_precision,
            layer_difference,
            corner_threshold,
            length_threshold,
            max_iterations,
            splice_threshold,
            path_precision
        ],
        outputs=[html, svg_output]
    )

# Launch the app
app.launch(debug=True)