import vtracer
import gradio as gr

from PIL import Image
from io import BytesIO
import base64

# Convert Image to Base64 
def im_2_b64(image):
    buff = BytesIO()
    image.save(buff, format="PNG")
    img_str = base64.b64encode(buff.getvalue())
    return img_str

# Save Base64 to png Image
def b64_2_img(data, out_file = './temp.png'):
    buff = BytesIO(base64.b64decode(data))
    img = Image.open(buff)
    img.save(out_file)
    return out_file


def png2svg(png_file,png_b64,svg_file,
            colormode,
            hierarchical,
            mode,
            filter_speckle,
            color_precision,
            layer_difference,
            corner_threshold,
            length_threshold,
            max_iterations,
            splice_threshold,
            path_precision):


    if len(png_b64):
        png_file = b64_2_img(png_b64.replace('data:image/png;base64,',''))
    
    params = {'colormode': colormode,
            'hierarchical': hierarchical,
            'mode':mode,
            'filter_speckle':filter_speckle,
            'color_precision':color_precision,
            'layer_difference':layer_difference,
            'corner_threshold':corner_threshold,
            'length_threshold':length_threshold,
            'max_iterations':max_iterations,
            'splice_threshold':splice_threshold,
            'path_precision':path_precision}

    vtracer.convert_image_to_svg_py(png_file, svg_file, **params)
    return svg_file

# Interface
with gr.Blocks() as interface:
    # Title
    gr.Markdown("# PNG to SVG Converter")
    # Heading
    gr.Markdown("Converts given png to svg using vtracer.")

    with gr.Accordion("File IO"):
        # file io
        png_file = gr.File(label = "Input file")
        png_b64 = gr.Textbox(label = "Input image b64")
        svg_file = gr.Textbox(label = "Output file", value="out.svg", interactive=True)

    # Parameters
    with gr.Accordion("Parameters"):
        colormode = gr.Dropdown(
            label = 'Color mode',
            choices = ['color','binary'],
            value = 'color',
            interactive=True)
        hierarchical = gr.Dropdown(
            label = 'Hierarchical',
            choices = ['stacked','cutout'],
            value = 'stacked',
            interactive=True)
        mode = gr.Dropdown(
            label = 'Curve fitting mode',
            choices = ['spline','polygon', 'pixel'],
            value = 'spline',
            interactive=True)
        filter_speckle = gr.Slider(
            label = 'Filter specles smaller than...',
            minimum = 0,
            maximum = 128,
            step = 1,
            value = 60,
            interactive=True)
        color_precision = gr.Slider(
            label = 'Color precision',
            minimum = 1,
            maximum = 8,
            step = 1,
            value = 8,
            interactive=True)
        layer_difference = gr.Slider(
            label = 'Gradient step size',
            minimum = 0,
            maximum = 128,
            step = 1,
            value = 16,
            interactive=True)
        corner_threshold = gr.Slider(
            label = 'Corner thershold',
            minimum = 0,
            maximum = 180,
            step = 1,
            value = 60,
            interactive=True)
        length_threshold = gr.Slider(
            label = 'Segment length',
            minimum = 3.5,
            maximum = 10,
            value = 10,
            interactive=True)
        max_iterations = gr.Slider(
            label = 'Max iterations',
            minimum = 3.5,
            maximum = 10,
            value = 10,
            interactive=True)
        splice_threshold = gr.Slider(
            label = 'Splice thereshold',
            minimum = 0,
            maximum = 180,
            value = 45,
            interactive=True)
        path_precision = gr.Slider(
            label = 'Path precision',
            minimum = 0,
            maximum = 16,
            step = 1,
            value = 8,
            interactive=True)

    # Convert
    convert_button = gr.Button("Convert")
    
    with gr.Accordion("Output SVG File"):
        out_file = gr.File()
        out_html = gr.HTML()
        
    # Event Bindings
    out_file.change((lambda x: f'<img src="file/{x}">'),out_file,out_html)
    convert_button.click(png2svg, 
        [
            png_file,
            png_b64,
            svg_file,
            colormode,
            hierarchical,
            mode,
            filter_speckle,
            color_precision,
            layer_difference,
            corner_threshold,
            length_threshold,
            max_iterations,
            splice_threshold,
            path_precision],
        [out_file])

                  
interface.launch()