import gradio as gr
import vtracer
import os
from PIL import Image

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"

    if image is None:
        return gr.HTML("Please upload an image."), None

    image.save(input_path)

    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)
    )

    try:
        with open(output_path, "r") as f:
            svg_content = f.read()
        return gr.HTML(f'<svg viewBox="0 0 {image.width} {image.height}">{svg_content}</svg>'), output_path
    except FileNotFoundError:
        return gr.HTML("SVG conversion failed."), None

def handle_color_mode(value):
    return value

examples = [
    "examples/11.jpg",
    "examples/02.jpg",
    "examples/03.jpg",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 960px;
    padding: 20px;
}
.generate-btn {
    background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%) !important;
    border: none !important;
    color: white !important;
    padding: 10px 20px !important;
    font-size: 16px !important;
    border-radius: 5px !important;
    cursor: pointer !important;
    transition: transform 0.2s, box-shadow 0.2s !important;
}
.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 5px 15px rgba(0,0,0,0.2) !important;
}
.gr-accordion {
    border: 1px solid #ddd !important;
    border-radius: 5px !important;
    margin-bottom: 10px !important;
}
.gr-accordion-title {
    background-color: #f9f9f9 !important;
    padding: 10px !important;
    cursor: pointer !important;
}
.gr-accordion-content {
    padding: 15px !important;
}
.gr-file {
    border: 1px dashed #aaa !important;
    padding: 10px !important;
    border-radius: 5px !important;
}
"""

with gr.Blocks(css=css) as app:
    with gr.Column(elem_id="col-container"):
        gr.HTML("""
            <div style="text-align: center; margin-bottom: 20px;">
                <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")
                with gr.Accordion("Advanced Settings", open=False):
                    with gr.Accordion("Clustering", open=False):
                        colormode = gr.Radio([("COLOR","color"),("B/W", "binary")], value="color", label="Color Mode", show_label=False)
                        filter_speckle = gr.Slider(0, 128, value=4, step=1, label="Filter Speckle", info="Cleaner")
                        color_precision = gr.Slider(1, 8, value=6, step=1, label="Color Precision", info="More accurate")
                        layer_difference = gr.Slider(0, 128, value=16, step=1, label="Gradient Step", info="Less layers")
                        hierarchical = gr.Radio([("STACKED","stacked"), ("CUTOUT","cutout")], value="stacked", label="Hierarchical Mode",show_label=False)
                    with gr.Accordion("Curve Fitting", open=False):
                        mode = gr.Radio([("SPLINE","spline"),("POLYGON", "polygon"), ("PIXEL","none")], value="spline", label="Mode", show_label=False)
                        corner_threshold = gr.Slider(0, 180, value=60, step=1, label="Corner Threshold", info="Smoother")
                        length_threshold = gr.Slider(3.5, 10, value=4.0, step=0.1, label="Segment Length", info ="More coarse")
                        splice_threshold = gr.Slider(0, 180, value=45, step=1, label="Splice Threshold", info="Less accurate")
                        max_iterations = gr.Slider(1, 20, value=10, step=1, label="Max Iterations", visible=False)
                        path_precision = gr.Slider(1, 10, value=3, step=1, label="Path Precision", visible=False)
                        output_text = gr.Textbox(label="Selected Mode", visible=False)
                with gr.Row():
                    clear_button = gr.Button("Clear")
                    convert_button = gr.Button("✨ Convert to SVG", variant='primary', elem_classes=["generate-btn"])

            with gr.Column():
                html = gr.HTML(label="SVG Output")
                svg_output = gr.File(label="Download SVG")
        
        gr.Examples(
            examples = examples,
            fn = convert_to_vector,
            inputs = [image_input],
            outputs = [html,svg_output],
            cache_examples=False,
            run_on_click = True
        )

    colormode.change(handle_color_mode, inputs=colormode, outputs=output_text)
    hierarchical.change(handle_color_mode, inputs=hierarchical, outputs=output_text)

    def update_mode_visibility(mode_val):
        is_spline_mode = mode_val == "spline"
        return (
            gr.update(interactive=is_spline_mode),
            gr.update(interactive=is_spline_mode),
            gr.update(interactive=is_spline_mode)
        )

    mode.change(
        update_mode_visibility,
        inputs=[mode],
        outputs=[corner_threshold, length_threshold, splice_threshold]
    )

    def clear_inputs():
        return gr.Image(value=None), gr.Radio(value="color"), gr.Radio(value="stacked"), gr.Radio(value="spline"), gr.Slider(value=4), gr.Slider(value=6), gr.Slider(value=16), gr.Slider(value=60), gr.Slider(value=4.0), gr.Slider(value=10), gr.Slider(value=45), gr.Slider(value=3)

    def update_colormode_visibility(colormode_val, color_prec, layer_diff):
        is_color_mode = colormode_val == "color"
        return (
            gr.update(interactive=is_color_mode),
            gr.update(interactive=is_color_mode),
            gr.update(visible=is_color_mode)