"""
G'MIC Filter Studio — Hugging Face Spaces
Runs the full modern G'MIC CLI (from apt) with Gradio UI.
Gives access to 600+ filters including community stdlib.
"""

import gradio as gr
import subprocess
import tempfile
import shutil
import os
import sys
from pathlib import Path
from PIL import Image

# ─── PATH TO GMIC ────────────────────────────────────────────────────────────
GMIC = shutil.which("gmic") or "/usr/bin/gmic"
STDLIB_UPDATE = os.path.expanduser("~/.config/gmic/update" + "xxx" + ".gmic")
# We'll resolve the actual update file path at startup

UPDATE_FETCHED = False


def _fetch_update_if_needed():
    """Download the G'MIC community filter definitions once at startup."""
    global UPDATE_FETCHED
    if UPDATE_FETCHED:
        return
    try:
        r = subprocess.run(
            [GMIC, "update"],
            capture_output=True, text=True, timeout=30
        )
        UPDATE_FETCHED = True
    except Exception as e:
        print(f"[gmic update] warning: {e}", file=sys.stderr)
        UPDATE_FETCHED = True   # continue even if offline


def _run_gmic(input_path: str, command: str, extra_prefix: str = "") -> tuple[str | None, str]:
    """
    Run a G'MIC command on input_path, return (output_path_or_None, log).
    extra_prefix: optional gmic args before the input file (e.g. extra script loading).
    """
    _fetch_update_if_needed()

    out_path = input_path.replace(".png", "_out.png")
    # Use a fresh temp dir so filenames are safe
    with tempfile.TemporaryDirectory() as td:
        inp = os.path.join(td, "input.png")
        out = os.path.join(td, "output.png")
        shutil.copy(input_path, inp)

        cmd = [GMIC]
        if extra_prefix:
            cmd += extra_prefix.split()
        cmd += [inp] + command.split() + ["-o", out]

        try:
            result = subprocess.run(
                cmd,
                capture_output=True, text=True, timeout=120
            )
            log = (result.stdout + "\n" + result.stderr).strip()

            if os.path.exists(out):
                final = input_path.replace(".png", "_gmic_result.png")
                shutil.copy(out, final)
                return final, log
            else:
                return None, log or "G'MIC produced no output file."
        except subprocess.TimeoutExpired:
            return None, "ERROR: G'MIC timed out (>120s). Try a smaller image or simpler filter."
        except Exception as e:
            return None, f"ERROR: {e}"


def _pil_to_temp(img: Image.Image) -> str:
    tmp = tempfile.mktemp(suffix=".png")
    img.save(tmp, "PNG")
    return tmp


# ─── FILTER DEFINITIONS ──────────────────────────────────────────────────────
# Each entry: (display_name, command_template, [(param_name, label, min, max, step, default)])
# {p_name} in command_template gets replaced with slider value.
# Grouped into tabs.

FILTER_GROUPS = {

    "📄 Document Repair": [
        ("Repair Scanned Document",
         "fx_repair_scanned_doc {brightness},{contrast},{gamma},{threshold},0",
         [("brightness", "Brightness", -100, 100, 1, 0),
          ("contrast",   "Contrast",   -100, 100, 1, 20),
          ("gamma",      "Gamma",      -100, 100, 1, 0),
          ("threshold",  "Threshold %", 0,  100, 1, 50)]),

        ("Clean Text (afre)",
         "afre_cleantext {threshold},{smoothness},{sharpen}",
         [("threshold",  "Threshold",  0, 100, 1, 50),
          ("smoothness", "Smoothness", 0,  10, 1,  2),
          ("sharpen",    "Sharpen",    0, 300, 5, 50)]),

        ("Denoise (Iain fast)",
         "iain_fast_denoise_p {strength},{patch},{search}",
         [("strength", "Strength",    1, 50,  1, 15),
          ("patch",    "Patch size",  1, 15,  2,  7),
          ("search",   "Search size", 5, 30,  1, 21)]),

        ("Remove Hot Pixels",
         "fx_remove_hotpixels {threshold},{mask}",
         [("threshold", "Threshold", 0,  30, 1,  5),
          ("mask",      "Mask size", 1,   5, 1,  2)]),

        ("Upscale 2×  (Scale2x)",
         "fx_scalenx 2", []),

        ("Deinterlace",
         "deinterlace 0", []),
    ],

    "🔧 Inpainting": [
        ("Inpaint – Patch-Based",
         "fx_inpaint_patch {patch},{overlap},{blend}",
         [("patch",   "Patch size",   5, 50, 1, 15),
          ("overlap", "Overlap",      0, 10, 1,  5),
          ("blend",   "Blend radius", 0, 10, 1,  1)]),

        ("Inpaint – Multi-Scale",
         "fx_inpaint_matchpatch {patch},{nb_scales}",
         [("patch",     "Patch size", 5, 60, 1, 11),
          ("nb_scales", "Scales",     1,  5, 1,  3)]),

        ("Fill Transparent Area",
         "fx_solidify_td {method},{nb_iters}",
         [("method",   "Method",     0, 3, 1, 0),
          ("nb_iters", "Iterations", 1, 5, 1, 1)]),
    ],

    "✨ Enhance / Restore": [
        ("Anisotropic Smooth (edge-preserving)",
         "smooth {amplitude},{sharpness},{anisotropy},{alpha},{sigma}",
         [("amplitude",  "Amplitude",  0, 100, 1,  40),
          ("sharpness",  "Sharpness",  0,   2, .1, .7),
          ("anisotropy", "Anisotropy", 0,   1, .05, .5),
          ("alpha",      "Alpha",      0,   4, .1,  .6),
          ("sigma",      "Sigma",      0,   4, .1, 1.1)]),

        ("Smooth Skin",
         "fx_smooth_skin 2,{tolerance},{smoothness},1,1,50,50,5,2,0.2,3,1,0.05,5,0,50,50",
         [("tolerance",  "Skin Tolerance", 0, 1, .05, .5),
          ("smoothness", "Smoothness",     0, 5, .1,  1.0)]),

        ("Retinex (local contrast)",
         "fx_retinex {variance},{part},{colorspace}",
         [("variance",    "Variance",    0, 600, 5, 200),
          ("part",        "Mix %",       0, 100, 5,  80),
          ("colorspace",  "Colorspace",  0,   3, 1,   1)]),

        ("DCP Dehaze",
         "jeje_dehaze {amount},{radius}",
         [("amount", "Amount", 0, 100, 1, 50),
          ("radius", "Radius", 1,  50, 1, 15)]),

        ("Sharpen (Texture)",
         "fx_sharpen_texture {amount},{scale}",
         [("amount", "Amount", 0, 300, 5,  100),
          ("scale",  "Scale",  0,   5, .1, 1.0)]),

        ("Unsharp Mask",
         "unsharp {std},{amount},{threshold}",
         [("std",       "Std dev",   .5, 10, .5,  2),
          ("amount",    "Amount",    0,   5, .1, 1.5),
          ("threshold", "Threshold", 0,  50,  1,   5)]),

        ("Normalize Local",
         "fx_normalize_tiles {size},{overlap},{normalize}",
         [("size",      "Tile size", 16, 256, 8,  48),
          ("overlap",   "Overlap",    0,  16, 1,   4),
          ("normalize", "Normalize",  0,   5, 1,   2)]),
    ],

    "🎨 Color & Tone": [
        ("Simulate Film",
         "fx_simulate_film {category},{index},{strength},{brightness},{contrast},{gamma},{hue},{saturation},{normalize},0,0,0",
         [("category",   "Category (0–9)",  0,   9, 1,  0),
          ("index",      "Film preset",     1,  10, 1,  1),
          ("strength",   "Strength",        0, 100, 5, 80),
          ("brightness", "Brightness",    -100,100, 1,  0),
          ("contrast",   "Contrast",      -100,100, 1,  0),
          ("gamma",      "Gamma",         -100,100, 1,  0),
          ("hue",        "Hue",            -180,180,5,  0),
          ("saturation", "Saturation",    -100,100, 1,  0),
          ("normalize",  "Normalize",       0,   3, 1,  0)]),

        ("Retro Fade",
         "fx_retrofade {strength},{colorfade},{vignette}",
         [("strength",  "Strength",   0, 100, 1, 50),
          ("colorfade", "Color Fade", 0, 100, 1, 30),
          ("vignette",  "Vignette",   0, 100, 1, 40)]),

        ("Color Temperature",
         "fx_tk_colortemp {temperature},{strength}",
         [("temperature", "Temperature (K)", 2000, 10000, 100, 6500),
          ("strength",    "Strength",           0,   100,   5,   80)]),

        ("Vibrance",
         "fx_vibrance {amount}",
         [("amount", "Amount", -100, 200, 1, 50)]),

        ("Color Grading",
         "jl_colorgrading {shadows},{midtones},{highlights},{saturation}",
         [("shadows",    "Shadows shift",    -50, 50, 1,  0),
          ("midtones",   "Midtones shift",   -50, 50, 1,  0),
          ("highlights", "Highlights shift", -50, 50, 1,  0),
          ("saturation", "Saturation",       -50, 50, 1,  0)]),

        ("Equalize HSV",
         "fx_hsv_equalizer {strength},{radius}",
         [("strength", "Strength", 0, 100, 1, 80),
          ("radius",   "Radius",   0,  50, 1, 20)]),
    ],

    "🖌️ Artistic": [
        ("Kuwahara Painting",
         "fx_kuwahara {radius},{sharpness}",
         [("radius",    "Radius",    1, 15, 1,  5),
          ("sharpness", "Sharpness", 0,  5, .1, 2.0)]),

        ("Brushify",
         "fx_brushify {size},{density},{opacity},{sharpness},{angle}",
         [("size",      "Brush size",    1, 30, 1,  6),
          ("density",   "Density",       0, 100,1, 50),
          ("opacity",   "Opacity",       0, 100,5, 80),
          ("sharpness", "Sharpness",     0,  10,.5, 2),
          ("angle",     "Angle (°)", -180, 180, 5,  0)]),

        ("Illustration Look",
         "fx_illustration_look {edges},{smoothness},{sharpness}",
         [("edges",      "Edges",      0, 100, 5,  50),
          ("smoothness", "Smoothness", 0,  10, .5,  3),
          ("sharpness",  "Sharpness",  0, 300, 10,100)]),

        ("Sketch (Pencil B&W)",
         "fx_sketchbw {amplitude},{edge_threshold},{smooth},{details}",
         [("amplitude",      "Amplitude",      0, 100, 1, 40),
          ("edge_threshold", "Edge threshold", 0, 100, 1, 20),
          ("smooth",         "Smoothness",     0,   5,.1, .5),
          ("details",        "Details",        0,   5,.1,  1)]),

        ("Cartoon",
         "cartoon {smooth},{sharp},{threshold},{quantize},{coeff}",
         [("smooth",    "Smooth",    0,  10, .5, 3),
          ("sharp",     "Sharp",     0, 300, 10,200),
          ("threshold", "Threshold", 0, 100,  1, 30),
          ("quantize",  "Quantize",  2,  12,  1,  8),
          ("coeff",     "Coeff",     0,   3,.1,  1)]),

        ("Vector Painting",
         "fx_vector_painting {detail}",
         [("detail", "Detail", 1, 15, 1, 7)]),

        ("Rodilius",
         "fx_rodilius {amplitude},{size},{smoothness},{sharpness}",
         [("amplitude",  "Amplitude",  0, 60, 1, 10),
          ("size",       "Size",       0, 20, 1,  6),
          ("smoothness", "Smoothness", 0,  4, .1, 1),
          ("sharpness",  "Sharpness",  0, 300,10,200)]),

        ("Poster Edges",
         "fx_poster_edges {edge_th},{smooth},{quantize}",
         [("edge_th",  "Edge threshold", 0, 100, 1, 20),
          ("smooth",   "Smoothness",     0,  10, .5, 1),
          ("quantize", "Quantize",       2,  12,  1,  8)]),
    ],

    "🔬 Analysis": [
        ("Edge Detection",
         "fx_edges {threshold}",
         [("threshold", "Threshold", 0, 100, 1, 15)]),

        ("Gradient Norm",
         "gradient_norm", []),

        ("Structure Tensor",
         "structure_tensors 0", []),

        ("Frequency Split (detail layer)",
         "-blur {radius} +[-2] -[-1] -n 0,255",
         [("radius", "Radius (blur size)", 1, 30, 1, 5)]),

        ("Local Variance Map",
         "+blur {radius} sqr[0] blur[1] {radius} sqr[1] - sqrt n 0,255",
         [("radius", "Radius", 1, 20, 1, 5)]),
    ],

}

# ─── CORE GRADIO FUNCTION ────────────────────────────────────────────────────

def apply_filter(input_image: Image.Image, filter_group: str,
                 filter_name: str, raw_command: str,
                 **slider_values) -> tuple[Image.Image | None, str]:
    """
    Main processing function called by Gradio.
    Either runs the raw_command directly, or looks up the preset.
    """
    if input_image is None:
        return None, "Please upload an image first."

    tmp_in = _pil_to_temp(input_image)

    # Decide command to run
    if raw_command.strip():
        cmd = raw_command.strip()
        log_prefix = f"[raw] {cmd}"
    else:
        # Find the filter
        filters = FILTER_GROUPS.get(filter_group, [])
        match = None
        for f in filters:
            if f[0] == filter_name:
                match = f
                break
        if match is None:
            return None, f"Filter '{filter_name}' not found in group '{filter_group}'."

        name, template, params = match
        vals = {}
        for pname, plabel, pmin, pmax, pstep, pdef in params:
            vals[pname] = slider_values.get(f"sl_{pname}", pdef)
        # Build command
        cmd = template
        for k, v in vals.items():
            cmd = cmd.replace("{" + k + "}", str(v))
        log_prefix = f"[{name}] {cmd}"

    out_path, log = _run_gmic(tmp_in, cmd)

    try:
        os.unlink(tmp_in)
    except Exception:
        pass

    if out_path and os.path.exists(out_path):
        result_img = Image.open(out_path)
        result_img = result_img.copy()
        try:
            os.unlink(out_path)
        except Exception:
            pass
        return result_img, log_prefix + "\n\n" + log
    else:
        return None, log_prefix + "\n\nFAILED:\n" + log


# ─── BUILD UI ────────────────────────────────────────────────────────────────

def build_ui():
    # All slider names across all filters
    all_params: dict[str, tuple] = {}
    for group_name, filters in FILTER_GROUPS.items():
        for fname, template, params in filters:
            for pname, plabel, pmin, pmax, pstep, pdef in params:
                key = f"sl_{pname}"
                if key not in all_params:
                    all_params[key] = (plabel, pmin, pmax, pstep, pdef)

    CSS = """
    #gmic-header { font-family: monospace; background: #0a0a0a; padding: 14px 20px;
                   border-bottom: 2px solid #d4ff00; }
    #gmic-header h1 { color: #d4ff00; font-size: 1.4em; margin: 0; letter-spacing: .08em; }
    #gmic-header p  { color: #888; font-size: .8em; margin: 4px 0 0; }
    .note { font-size: .8em; color: #888; font-family: monospace; }
    """

    with gr.Blocks(css=CSS, title="G'MIC Filter Studio") as demo:

        gr.HTML("""
        <div id="gmic-header">
          <h1>G'MIC FILTER STUDIO</h1>
          <p>Full G'MIC stdlib + community filters · runs on server · no size limit</p>
        </div>
        """)

        with gr.Row():

            # ── LEFT: inputs ──
            with gr.Column(scale=1, min_width=300):
                input_img = gr.Image(
                    type="pil", label="Input Image"
                )

                gr.HTML("<div class='note'>Upload any size. Server-side processing on CPU.</div>")

                group_dd = gr.Dropdown(
                    label="Filter Category",
                    choices=list(FILTER_GROUPS.keys()),
                    value=list(FILTER_GROUPS.keys())[0]
                )

                filter_dd = gr.Dropdown(
                    label="Filter",
                    choices=[f[0] for f in FILTER_GROUPS[list(FILTER_GROUPS.keys())[0]]],
                    value=FILTER_GROUPS[list(FILTER_GROUPS.keys())[0]][0][0]
                )

                # All sliders, shown/hidden via JS class below
                slider_components: dict[str, gr.Slider] = {}
                with gr.Group(visible=True) as param_group:
                    gr.Markdown("**Parameters**")
                    for key, (plabel, pmin, pmax, pstep, pdef) in all_params.items():
                        sl = gr.Slider(
                            minimum=pmin, maximum=pmax, step=pstep, value=pdef,
                            label=plabel, visible=False, elem_id=key
                        )
                        slider_components[key] = sl

                raw_cmd = gr.Textbox(
                    label="Raw G'MIC command (overrides preset if non-empty)",
                    placeholder="-blur 3 -edges 5\nfx_repair_scanned_doc 0,20,0,50,0\n...",
                    lines=3,
                    info="Applied directly to the input image. Ctrl+Enter to run."
                )

                apply_btn = gr.Button("▶  Apply Filter", variant="primary")

            # ── RIGHT: outputs ──
            with gr.Column(scale=1, min_width=300):
                output_img = gr.Image(type="pil", label="Result", interactive=False)
                log_box    = gr.Textbox(label="G'MIC log", lines=6, interactive=False)

                with gr.Accordion("📖 Filter Reference", open=False):
                    gr.Markdown("""
### Filters unique to G'MIC (not in ImageJ.js / Photopea)

| Category | Filter | G'MIC command |
|---|---|---|
| Document | **Repair Scanned Document** | `fx_repair_scanned_doc` |
| Inpainting | Patch-based inpaint | `fx_inpaint_patch` |
| Inpainting | Multi-scale inpaint | `fx_inpaint_matchpatch` |
| Smoothing | Anisotropic diffusion | `smooth` |
| Skin | Smooth Skin | `fx_smooth_skin` |
| Tone | Retinex | `fx_retinex` |
| Tone | DCP Dehaze | `jeje_dehaze` |
| Film | Simulate Film (600+ LUTs) | `fx_simulate_film` |
| Artistic | Brushify | `fx_brushify` |
| Artistic | Kuwahara | `fx_kuwahara` |
| B&W | Engrave / Filaments | `fx_engrave` |
| Analysis | Frequency split | pipeline |
| Upscale | Scale2x | `fx_scalenx` |
| Denoise | Iain Fast Denoise | `iain_fast_denoise_p` |

**Tip — find any filter's CLI args:**
In GIMP G'MIC plugin, set *Output Messages → Verbose (layer name)*, apply the filter,
and the layer name shows the exact CLI command. Or:
```
gmic echo '$${fx_some_filter}'
```
""")

        # ── Dynamics ──────────────────────────────────────────────

        def update_filter_list(group_name):
            filters = FILTER_GROUPS.get(group_name, [])
            choices = [f[0] for f in filters]
            val = choices[0] if choices else None
            return gr.update(choices=choices, value=val)

        def update_sliders(group_name, filter_name):
            # Find the matching filter
            filters = FILTER_GROUPS.get(group_name, [])
            match = next((f for f in filters if f[0] == filter_name), None)
            updates = {}
            if match:
                _, _, params = match
                active_pnames = {f"sl_{p[0]}" for p in params}
                for key in slider_components:
                    if key in active_pnames:
                        param = next(p for p in params if f"sl_{p[0]}" == key)
                        updates[slider_components[key]] = gr.update(
                            visible=True, label=param[1],
                            minimum=param[2], maximum=param[3],
                            step=param[4], value=param[5]
                        )
                    else:
                        updates[slider_components[key]] = gr.update(visible=False)
            else:
                for sl in slider_components.values():
                    updates[sl] = gr.update(visible=False)
            return list(updates.values())

        group_dd.change(
            update_filter_list,
            inputs=[group_dd],
            outputs=[filter_dd]
        )
        filter_dd.change(
            update_sliders,
            inputs=[group_dd, filter_dd],
            outputs=list(slider_components.values())
        )
        # Initialize slider visibility on page load
        demo.load(
            update_sliders,
            inputs=[group_dd, filter_dd],
            outputs=list(slider_components.values())
        )

        # Wire apply button
        all_slider_comps = list(slider_components.values())
        all_slider_keys  = list(slider_components.keys())

        def on_apply(input_image, group_name, filter_name, raw_command, *slider_vals):
            sv = {all_slider_keys[i]: slider_vals[i] for i in range(len(slider_vals))}
            return apply_filter(input_image, group_name, filter_name, raw_command, **sv)

        apply_btn.click(
            on_apply,
            inputs=[input_img, group_dd, filter_dd, raw_cmd] + all_slider_comps,
            outputs=[output_img, log_box]
        )

        # Also trigger on Ctrl+Enter in raw_cmd
        raw_cmd.submit(
            on_apply,
            inputs=[input_img, group_dd, filter_dd, raw_cmd] + all_slider_comps,
            outputs=[output_img, log_box]
        )

    return demo


# ─── STARTUP ────────────────────────────────────────────────────────────────

# Prefetch gmic update file in background on cold start
import threading
threading.Thread(target=_fetch_update_if_needed, daemon=True).start()

demo = build_ui()

if __name__ == "__main__":
    demo.launch()