import cv2
import numpy as np
import gradio as gr
import tempfile


def process(image, mode, k):
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    clahe = cv2.createCLAHE(2.0, (8, 8))
    gray = clahe.apply(gray)

    k = max(3, int(k))
    if k % 2 == 0:
        k += 1

    if mode == "Sketch":
        inv = 255 - gray
        blur = cv2.GaussianBlur(inv, (k, k), 0)
        out = cv2.divide(gray, 255 - blur, scale=256)

    elif mode == "Embroidery":
        blur = cv2.GaussianBlur(gray, (k, k), 0)

        out = cv2.adaptiveThreshold(
            blur,
            255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY,
            11,
            2
        )

    elif mode == "Cartoon":
        color = cv2.bilateralFilter(img, 9, 75, 75)

        gray_c = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        edges = cv2.adaptiveThreshold(
            gray_c,
            255,
            cv2.ADAPTIVE_THRESH_MEAN_C,
            cv2.THRESH_BINARY,
            9,
            2
        )

        edges = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
        out = cv2.bitwise_and(color, edges)

    else:
        out = gray

    # Create temp file for download
    temp_file = tempfile.NamedTemporaryFile(
        suffix=".png",
        delete=False
    )

    cv2.imwrite(temp_file.name, out)

    # Preview image for Gradio
    if len(out.shape) == 2:
        preview = out
    else:
        preview = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)

    return preview, temp_file.name


demo = gr.Interface(
    fn=process,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Dropdown(
            ["Sketch", "Embroidery", "Cartoon"],
            value="Sketch",
            label="Mode"
        ),
        gr.Slider(
            minimum=3,
            maximum=50,
            value=9,
            step=1,
            label="Strength"
        )
    ],
    outputs=[
        gr.Image(label="Preview"),
        gr.File(label="Download PNG")
    ],
    title= "Sketch Converter Studio",
    description="Convert images into sketch, embroidery line art, or cartoon styles."
)

demo.launch()