# normalmap_to_directx_filtered.py
# Drag & drop image files or folders onto the EXE.
# Default: FORCE convert to DirectX (−Y) by flipping green.
# Suffix filter: only process files whose *basename* ends with one of: _n, _nrm, _normal (case-insensitive).
# Flags:
#   --detect              -> flip only if image looks OpenGL (+Y)
#   --all                 -> ignore suffix filter; process all supported images
#   --suffixes "a,b,c"    -> comma-separated list of suffixes (without extensions)

import sys, os
from PIL import Image
import numpy as np

SUPPORTED_EXTS = {".png", ".tga", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"}
OUT_SUBFOLDER = "DirectX_Converted"
DEFAULT_SUFFIXES = ["_n", "_nrm", "_normal"]  # case-insensitive

def is_image(p):
    return os.path.splitext(p)[1].lower() in SUPPORTED_EXTS

def iter_inputs(paths):
    for p in paths:
        if os.path.isdir(p):
            for root, _, files in os.walk(p):
                for f in files:
                    fp = os.path.join(root, f)
                    if is_image(fp):
                        yield fp
        else:
            if is_image(p):
                yield p

def has_normal_suffix(path, suffixes):
    stem = os.path.splitext(os.path.basename(path))[0].lower()
    return any(stem.endswith(suf.lower()) for suf in suffixes)

def flip_green(img_rgba):
    r, g, b, a = img_rgba.split()
    g = g.point(lambda i: 255 - i)
    return Image.merge("RGBA", (r, g, b, a))

def analyze_mean_y(img_rgba):
    _, g, b, _ = img_rgba.split()
    g_np = np.array(g, dtype=np.float32)
    b_np = np.array(b, dtype=np.float32)
    y = (g_np / 255.0) * 2.0 - 1.0
    flat_mask = b_np > 240
    y_use = y[~flat_mask] if (~flat_mask).any() else y
    return float(y_use.mean())

def ensure_directx(img_rgba, detect_mode: bool):
    if not detect_mode:
        return flip_green(img_rgba), True, None  # forced
    mean_y = analyze_mean_y(img_rgba)
    if mean_y >= 0.0:
        return flip_green(img_rgba), True, mean_y
    else:
        return img_rgba, False, mean_y

def output_path(src_path):
    folder, base = os.path.split(src_path)
    out_dir = os.path.join(folder, OUT_SUBFOLDER)
    os.makedirs(out_dir, exist_ok=True)
    return os.path.join(out_dir, base)

def save_preserving_format(out_img_rgba, src_path, had_alpha):
    _, ext = os.path.splitext(src_path)
    ext = ext.lower()

    if not had_alpha or ext in {".jpg", ".jpeg", ".bmp"}:
        out_img = out_img_rgba.convert("RGB")
    else:
        out_img = out_img_rgba

    dst = output_path(src_path)
    save_kwargs, fmt = {}, None
    if ext in {".jpg", ".jpeg"}:
        save_kwargs["quality"] = 95
        fmt = "JPEG"
    elif ext == ".png":
        fmt = "PNG"
    elif ext == ".tga":
        fmt = "TGA"
    elif ext in {".tif", ".tiff"}:
        fmt = "TIFF"
    elif ext == ".bmp":
        fmt = "BMP"
    else:
        dst = os.path.splitext(dst)[0] + ".png"
        fmt = "PNG"

    out_img.save(dst, format=fmt, **save_kwargs)
    return dst

def process_one(path, detect_mode: bool):
    try:
        src = Image.open(path)
        had_alpha = src.mode in ("LA", "RGBA", "PA")
        img = src.convert("RGBA")

        out_img, flipped, mean_y = ensure_directx(img, detect_mode)
        dst = save_preserving_format(out_img, path, had_alpha)

        if detect_mode:
            status = "flipped to DirectX (−Y)" if flipped else "already DirectX (−Y)"
            extra = f"  meanY={mean_y:+.4f}"
        else:
            status = "FORCED flip -> DirectX (−Y)"
            extra = ""
        print(f"[OK] {path}\n     {status}{extra}\n     -> {dst}")
    except Exception as e:
        print(f"[ERR] {path} :: {e}")

def parse_args(argv):
    detect_mode = False
    process_all = False
    suffixes = DEFAULT_SUFFIXES[:]
    paths = []
    it = iter(argv)
    for a in it:
        if a == "--detect":
            detect_mode = True
        elif a == "--all":
            process_all = True
        elif a == "--suffixes":
            try:
                raw = next(it)
                suffixes = [s.strip() for s in raw.split(",") if s.strip()]
            except StopIteration:
                print("[WARN] --suffixes expects a quoted comma-separated list; using defaults.")
        else:
            paths.append(a)
    return detect_mode, process_all, suffixes, paths

def main():
    detect_mode, process_all, suffixes, args = parse_args(sys.argv[1:])

    if not args:
        print("Drag and drop image files or folders onto this EXE.")
        print("Options: --detect  --all  --suffixes \"_n,_nrm,_normal\"")
        return  # auto-exit

    files = list(iter_inputs(args))
    if not files:
        print("No supported images found.")
        return

    mode_desc = "DETECT mode (flip only if +Y detected)" if detect_mode else "FORCE mode (flip everything)"
    filt_desc = "NO suffix filter (--all)" if process_all else f"Suffix filter: {', '.join(suffixes)}"
    print(f"{mode_desc}\n{filt_desc}\nFound {len(files)} file(s) before filtering.\n")

    count_total = 0
    count_skipped = 0
    for p in files:
        if not process_all and not has_normal_suffix(p, suffixes):
            print(f"[SKIP] {p} (name lacks normal-map suffix)")
            count_skipped += 1
            continue
        count_total += 1
        process_one(p, detect_mode)

    print(f"\nProcessed: {count_total}, Skipped: {count_skipped}")

if __name__ == "__main__":
    main()