import gradio as gr
from PIL import Image, ImageFilter
import io, os, re
import numpy as np
import vtracer

def remove_background(arr):
    """Remove fundo detectando cor dos cantos."""
    h, w = arr.shape[:2]
    corners = []
    for cy, cx in [(0,0),(0,w-1),(h-1,0),(h-1,w-1),(h//2,0),(h//2,w-1),(0,w//2),(h-1,w//2)]:
        corners.append(arr[cy,cx,:3])
    bg_color = np.median(corners, axis=0)
    bg_dist = np.sqrt(((arr[:,:,:3].astype(np.float32)-bg_color)**2).sum(axis=2))
    arr[:,:,3] = np.where(bg_dist < 30, 0, arr[:,:,3])
    return arr, bg_color

def get_dominant_colors(arr, n_colors):
    """Extrai cores dominantes via quantizacao."""
    opaque = arr[:,:,3] > 128
    pixels = arr[:,:,:3][opaque].astype(np.uint8)
    if len(pixels) < n_colors:
        return []
    sample = Image.fromarray(pixels.reshape(-1,1,3), 'RGB')
    q = sample.quantize(colors=int(n_colors), method=Image.Quantize.MEDIANCUT)
    pal = q.getpalette()
    colors = []
    for i in range(int(n_colors)):
        c = np.array([pal[i*3],pal[i*3+1],pal[i*3+2]], dtype=np.float32)
        lum = 0.299*c[0]/255 + 0.587*c[1]/255 + 0.114*c[2]/255
        if lum > 0.05:  # ignora preto puro
            colors.append(c)
    return colors

def snap_to_dominant(arr, dominant, threshold=35):
    """Forca cada pixel para a cor dominante mais proxima."""
    if not dominant:
        return arr
    h, w = arr.shape[:2]
    dom_arr = np.array(dominant, dtype=np.float32)
    flat = arr[:,:,:3].reshape(-1,3).astype(np.float32)
    opaque_flat = arr[:,:,3].reshape(-1) > 128

    dists = np.sqrt(((flat[:,np.newaxis,:]-dom_arr[np.newaxis,:,:])**2).sum(axis=2))
    best_idx = np.argmin(dists, axis=1)
    best_dist = dists[np.arange(len(flat)), best_idx]
    best_color = dom_arr[best_idx]

    # Snap agressivo: todos os pixels opacos vao para a cor dominante
    snap = opaque_flat  # snap em TODOS os pixels opacos
    flat[snap] = best_color[snap]
    arr[:,:,:3] = flat.reshape(h,w,3)
    return arr

def clean_borders(arr):
    """Remove pixels de borda com cor misturada usando erosao + expansao."""
    try:
        import cv2
        alpha = arr[:,:,3].astype(np.uint8)
        k3 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3))
        k5 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))

        # Erode para remover pixels de borda misturados
        alpha_clean = cv2.erode(alpha, k3, iterations=2)
        # Fecha para recuperar forma interna
        alpha_clean = cv2.morphologyEx(alpha_clean, cv2.MORPH_CLOSE, k5, iterations=2)
        arr[:,:,3] = np.where(alpha_clean > 127, 255, 0)
    except ImportError:
        arr[:,:,3] = np.where(arr[:,:,3] > 128, 255, 0)
    return arr

def preprocess(pil_img, shadow_strength, n_colors, upscale):
    img = pil_img.convert('RGBA')
    w, h = img.size

    MAX, MIN = 1000, 400
    s = 1.0
    if max(w,h) > MAX: s = MAX/max(w,h)
    elif max(w,h) < MIN: s = MIN/max(w,h)
    if s != 1.0:
        img = img.resize((int(w*s), int(h*s)), Image.LANCZOS)

    arr = np.array(img, dtype=np.float32)

    # 1. Remove fundo automaticamente
    arr, bg_color = remove_background(arr)

    r,g,b,a = arr[:,:,0],arr[:,:,1],arr[:,:,2],arr[:,:,3]
    opaque = a > 128

    # 2. Remove sombras por saturacao
    if shadow_strength > 0:
        cmax = np.maximum(np.maximum(r,g),b)
        cmin = np.minimum(np.minimum(r,g),b)
        sat = np.where(cmax>0,(cmax-cmin)/cmax,0)
        lum = (0.299*r+0.587*g+0.114*b)/255
        shadow = (sat < shadow_strength*0.35) & opaque
        for ch in range(3):
            arr[:,:,ch] = np.where(shadow, np.where(lum>0.5,255,0), arr[:,:,ch])

    # 3. Alpha threshold
    arr[:,:,3] = np.where(arr[:,:,3]>128, 255, 0)

    # 4. Detecta cores dominantes
    dominant = get_dominant_colors(arr.astype(np.uint8), int(n_colors))

    # 5. Snap AGRESSIVO para cor dominante — elimina todas as variações
    arr = snap_to_dominant(arr, dominant, threshold=40)

    # 6. Limpa bordas com erosao
    arr = clean_borders(arr)

    img = Image.fromarray(arr.astype(np.uint8), 'RGBA')

    # 7. Upscale
    if upscale > 1:
        nw, nh = int(img.size[0]*upscale), int(img.size[1]*upscale)
        img = img.resize((nw, nh), Image.LANCZOS)
        arr2 = np.array(img, dtype=np.float32)
        try:
            import cv2
            alpha2 = arr2[:,:,3].astype(np.uint8)
            k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3))
            alpha2 = cv2.erode(alpha2, k, iterations=1)
            alpha2 = cv2.dilate(alpha2, k, iterations=1)
            arr2[:,:,3] = np.where(alpha2>128, 255, 0)
        except ImportError:
            arr2[:,:,3] = np.where(arr2[:,:,3]>128, 255, 0)

        # Snap novamente após upscale para eliminar pixels intermediários
        arr2 = snap_to_dominant(arr2, dominant, threshold=40)
        img = Image.fromarray(arr2.astype(np.uint8), 'RGBA')

    return img


def vectorize_image(
    input_image, shadow_strength, n_colors, upscale,
    filter_speckle, color_precision, corner_threshold,
    length_threshold, path_precision,
):
    try:
        if input_image is None:
            raise gr.Error("Nenhuma imagem enviada")

        pil_img = input_image.convert('RGBA')
        w, h = pil_img.size
        print(f"Entrada: {w}x{h}")
        if w>2000 or h>2000:
            raise gr.Error("Imagem muito grande. Maximo: 2000x2000 px.")

        pil_img = preprocess(pil_img, shadow_strength, n_colors, float(upscale))
        pw, ph = pil_img.size
        print(f"Processada: {pw}x{ph}")

        buf = io.BytesIO()
        pil_img.save(buf, format='PNG')

        svg_str = vtracer.convert_raw_image_to_svg(
            buf.getvalue(),
            img_format='png',
            colormode='color',
            hierarchical='stacked',
            mode='spline',
            filter_speckle=int(filter_speckle),
            color_precision=int(color_precision),
            layer_difference=16,
            corner_threshold=int(corner_threshold),
            length_threshold=float(length_threshold),
            max_iterations=10,
            splice_threshold=45,
            path_precision=int(path_precision),
        )

        if 'viewBox' not in svg_str:
            svg_str = svg_str.replace('<svg ', f'<svg viewBox="0 0 {pw} {ph}" ', 1)
        svg_str = re.sub(r'(<svg[^>]*?)width="[^"]*"', r'\1width="100%"', svg_str)
        svg_str = re.sub(r'(<svg[^>]*?)height="[^"]*"', r'\1height="100%"', svg_str)

        svg_path = f"/tmp/result_{os.getpid()}.svg"
        with open(svg_path,'w',encoding='utf-8') as f:
            f.write(svg_str)

        png_path = None
        try:
            import cairosvg
            png_bytes = cairosvg.svg2png(bytestring=svg_str.encode(), scale=1.0)
            png_path = f"/tmp/result_{os.getpid()}.png"
            with open(png_path,'wb') as f: f.write(png_bytes)
        except Exception as e:
            print(f"cairosvg: {e}")

        print(f"SVG: {len(svg_str)} chars")
        return svg_path, png_path, svg_str

    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(str(e))


with gr.Blocks(title="Iluminados Vectorizer") as demo:
    gr.Markdown("## Iluminados Vectorizer - VTracer API")

    with gr.Row():
        with gr.Column():
            inp = gr.Image(label="Imagem de entrada", type="pil")

            with gr.Accordion("Pre-processamento", open=True):
                shadow_strength = gr.Slider(0,1.0,value=0.6, step=0.05,label="Remover sombras")
                n_colors        = gr.Slider(2,32, value=16,  step=1,   label="Numero de cores dominantes")
                upscale         = gr.Slider(1,4,  value=2,   step=0.5, label="Upscale")

            with gr.Accordion("Vetorizacao", open=False):
                filter_speckle   = gr.Slider(1,100,value=4,  step=1,   label="Filtro de fragmentos")
                color_precision  = gr.Slider(1,8,  value=6,  step=1,   label="Precisao de cor")
                corner_threshold = gr.Slider(1,180,value=60, step=1,   label="Limiar de canto")
                length_threshold = gr.Slider(0.5,10,value=4.0,step=0.5,label="Comprimento minimo")
                path_precision   = gr.Slider(1,10, value=8,  step=1,   label="Precisao do caminho")

            btn = gr.Button("Vetorizar", variant="primary")

        with gr.Column():
            out_svg  = gr.File(label="SVG para download")
            out_png  = gr.File(label="PNG vetorizado")
            out_code = gr.Code(label="SVG codigo", language="html", lines=10)

    btn.click(
        fn=vectorize_image,
        inputs=[inp, shadow_strength, n_colors, upscale,
                filter_speckle, color_precision,
                corner_threshold, length_threshold, path_precision],
        outputs=[out_svg, out_png, out_code]
    )

demo.launch()