#!/usr/bin/env python3
"""Combine per-material Substance exports with a Blender Material ID map.

The companion Blender extension writes a PNG and JSON legend.  This tool uses
those exact ID colors as masks, discovers material names in Substance Painter
filenames, and combines matching channels without trusting texture alpha.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Sequence

try:
    from PIL import Image, ImageChops
except ImportError:  # pragma: no cover - exercised by the friendly CLI error.
    Image = None
    ImageChops = None


SUPPORTED_SUFFIXES = {
    ".bmp",
    ".jpeg",
    ".jpg",
    ".png",
    ".tga",
    ".tif",
    ".tiff",
    ".webp",
}
MATERIAL_TOKEN = "{material}"


class CombineError(RuntimeError):
    """A user-correctable texture discovery or composition error."""


@dataclass(frozen=True)
class LegendMaterial:
    slot: int
    name: str
    rgb: tuple[int, int, int]
    filename_name: str


@dataclass(frozen=True)
class TextureCandidate:
    path: Path
    material: LegendMaterial
    relative_parent: Path
    stem_pattern: str
    suffix: str


@dataclass
class TextureGroup:
    relative_parent: Path
    stem_pattern: str
    suffix: str
    textures: dict[str, Path]

    @property
    def output_name(self) -> str:
        stem = self.stem_pattern.replace(MATERIAL_TOKEN, "")
        stem = re.sub(r"[\s_.-]+", "_", stem).strip("_")
        return f"{stem or 'Combined'}{self.suffix.lower()}"

    @property
    def output_parent(self) -> Path:
        parts = []
        for part in self.relative_parent.parts:
            cleaned = part.replace(MATERIAL_TOKEN, "")
            cleaned = re.sub(r"[\s_.-]+", "_", cleaned).strip("_")
            if cleaned:
                parts.append(cleaned)
        return Path(*parts) if parts else Path(".")

    @property
    def description(self) -> str:
        parent = "" if self.relative_parent == Path(".") else f"{self.relative_parent}/"
        return f"{parent}{self.stem_pattern}{self.suffix}"


def _require_pillow() -> None:
    if Image is None or ImageChops is None:
        raise CombineError(
            "Pillow is required. Install it with "
            "'python -m pip install -r tools/requirements.txt'."
        )


def _parse_aliases(values: Sequence[str]) -> dict[str, str]:
    aliases: dict[str, str] = {}
    for value in values:
        if "=" not in value:
            raise CombineError(
                f"Invalid material alias {value!r}; expected 'LEGEND NAME=FILENAME NAME'."
            )
        legend_name, filename_name = (part.strip() for part in value.split("=", 1))
        if not legend_name or not filename_name:
            raise CombineError(f"Invalid material alias {value!r}; both names are required.")
        aliases[legend_name.casefold()] = filename_name
    return aliases


def _rgb8(value: Any, material_name: str) -> tuple[int, int, int]:
    if not isinstance(value, list) or len(value) < 3:
        raise CombineError(f"Legend material {material_name!r} has no valid RGB triplet.")
    try:
        components = tuple(float(component) for component in value[:3])
    except (TypeError, ValueError) as exc:
        raise CombineError(
            f"Legend material {material_name!r} has a non-numeric RGB value."
        ) from exc
    if any(component < 0.0 or component > 1.0 for component in components):
        raise CombineError(
            f"Legend material {material_name!r} has RGB values outside 0.0–1.0."
        )
    return tuple(round(component * 255.0) for component in components)


def load_legend(path: Path, aliases: dict[str, str] | None = None) -> list[LegendMaterial]:
    aliases = aliases or {}
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except OSError as exc:
        raise CombineError(f"Could not read legend {path}: {exc}") from exc
    except json.JSONDecodeError as exc:
        raise CombineError(f"Legend {path} is not valid JSON: {exc}") from exc

    entries = payload.get("materials") if isinstance(payload, dict) else None
    if not isinstance(entries, list) or not entries:
        raise CombineError(f"Legend {path} does not contain a non-empty 'materials' list.")

    materials: list[LegendMaterial] = []
    seen_names: set[str] = set()
    seen_colors: set[tuple[int, int, int]] = set()
    for fallback_slot, entry in enumerate(entries, start=1):
        if not isinstance(entry, dict):
            raise CombineError(f"Legend entry {fallback_slot} is not an object.")
        name = str(entry.get("material", "")).strip()
        if not name:
            raise CombineError(f"Legend entry {fallback_slot} has no material name.")
        folded_name = name.casefold()
        if folded_name in seen_names:
            raise CombineError(f"Legend contains duplicate material name {name!r}.")
        rgb = _rgb8(entry.get("rgb"), name)
        if rgb in seen_colors:
            raise CombineError(
                f"Legend color {rgb} is shared by multiple materials; masks would be ambiguous."
            )
        try:
            slot = int(entry.get("slot", fallback_slot))
        except (TypeError, ValueError) as exc:
            raise CombineError(f"Legend material {name!r} has an invalid slot.") from exc
        materials.append(
            LegendMaterial(
                slot=slot,
                name=name,
                rgb=rgb,
                filename_name=aliases.get(folded_name, name),
            )
        )
        seen_names.add(folded_name)
        seen_colors.add(rgb)
    return materials


def _name_pattern(name: str) -> re.Pattern[str]:
    tokens = re.findall(r"[A-Za-z0-9]+", name)
    if not tokens:
        raise CombineError(f"Material filename name {name!r} has no letters or digits.")
    expression = r"[\s_.-]*".join(re.escape(token) for token in tokens)
    return re.compile(
        rf"(?<![A-Za-z0-9]){expression}(?![A-Za-z0-9])",
        flags=re.IGNORECASE,
    )


def _match_candidate(
    path: Path,
    input_dir: Path,
    materials: Sequence[LegendMaterial],
) -> TextureCandidate | None:
    relative = path.relative_to(input_dir)
    searchable = [("stem", relative.stem)] + [
        (f"parent:{index}", part)
        for index, part in enumerate(relative.parent.parts)
    ]
    matches: list[
        tuple[int, int, LegendMaterial, re.Match[str], str]
    ] = []
    for material in materials:
        pattern = _name_pattern(material.filename_name)
        for location, text in searchable:
            match = pattern.search(text)
            if match:
                matches.append(
                    (
                        match.end() - match.start(),
                        len(material.filename_name),
                        material,
                        match,
                        location,
                    )
                )
    if not matches:
        return None

    _, _, material, match, location = max(
        matches,
        key=lambda item: (item[0], item[1], item[4] == "stem"),
    )
    stem_pattern = relative.stem
    parent_parts = list(relative.parent.parts)
    if location == "stem":
        stem_pattern = (
            f"{stem_pattern[:match.start()]}{MATERIAL_TOKEN}{stem_pattern[match.end():]}"
        )
    else:
        parent_index = int(location.split(":", 1)[1])
        parent_part = parent_parts[parent_index]
        parent_parts[parent_index] = (
            f"{parent_part[:match.start()]}{MATERIAL_TOKEN}{parent_part[match.end():]}"
        )
    relative_parent = Path(*parent_parts) if parent_parts else Path(".")
    return TextureCandidate(
        path=path,
        material=material,
        relative_parent=relative_parent,
        stem_pattern=stem_pattern,
        suffix=path.suffix,
    )


def discover_groups(
    input_dir: Path,
    materials: Sequence[LegendMaterial],
    recursive: bool = False,
    excluded_paths: Iterable[Path] = (),
) -> tuple[list[TextureGroup], list[Path]]:
    excluded = {path.resolve() for path in excluded_paths}
    iterator = input_dir.rglob("*") if recursive else input_dir.glob("*")
    files = sorted(
        path
        for path in iterator
        if path.is_file()
        and path.suffix.lower() in SUPPORTED_SUFFIXES
        and path.resolve() not in excluded
    )

    grouped: dict[tuple[Path, str, str], TextureGroup] = {}
    unmatched: list[Path] = []
    for path in files:
        candidate = _match_candidate(path, input_dir, materials)
        if candidate is None:
            unmatched.append(path)
            continue
        key = (
            candidate.relative_parent,
            candidate.stem_pattern.casefold(),
            candidate.suffix.casefold(),
        )
        group = grouped.setdefault(
            key,
            TextureGroup(
                relative_parent=candidate.relative_parent,
                stem_pattern=candidate.stem_pattern,
                suffix=candidate.suffix,
                textures={},
            ),
        )
        material_key = candidate.material.name.casefold()
        if material_key in group.textures:
            raise CombineError(
                f"Both {group.textures[material_key]} and {candidate.path} match material "
                f"{candidate.material.name!r} in group {group.description!r}."
            )
        group.textures[material_key] = candidate.path

    return sorted(grouped.values(), key=lambda group: group.description.casefold()), unmatched


def _mask_for_color(id_rgb: Any, rgb: tuple[int, int, int], tolerance: int) -> Any:
    channel_masks = []
    for channel, expected in zip(id_rgb.split(), rgb, strict=True):
        lookup = [255 if abs(value - expected) <= tolerance else 0 for value in range(256)]
        channel_masks.append(channel.point(lookup, mode="L"))
    mask = ImageChops.multiply(
        ImageChops.multiply(channel_masks[0], channel_masks[1]),
        channel_masks[2],
    )
    # One bit is sufficient for exact material membership and keeps 4K/8K
    # projects from retaining a full byte per material per pixel.
    return mask.convert("1")


def _common_mode(images: Sequence[Any]) -> str:
    modes = {image.mode for image in images}
    if len(modes) == 1:
        mode = next(iter(modes))
        if mode in {"1", "L", "LA", "RGB", "RGBA", "I", "F"}:
            return mode
    if any("A" in mode or mode == "P" for mode in modes):
        return "RGBA"
    if modes <= {"1", "L", "I", "F"}:
        return "L"
    return "RGB"


def combine_group(
    group: TextureGroup,
    materials: Sequence[LegendMaterial],
    masks: dict[str, Any],
    expected_size: tuple[int, int],
    output_path: Path,
    overwrite: bool = False,
) -> list[str]:
    if output_path.exists() and not overwrite:
        raise CombineError(f"Output already exists: {output_path} (use --overwrite).")

    opened: dict[str, Any] = {}
    try:
        for material in materials:
            path = group.textures.get(material.name.casefold())
            if path is None:
                continue
            image = Image.open(path)
            image.load()
            if image.size != expected_size:
                raise CombineError(
                    f"{path} is {image.size[0]}x{image.size[1]}, but the ID map is "
                    f"{expected_size[0]}x{expected_size[1]}."
                )
            opened[material.name.casefold()] = image

        if not opened:
            raise CombineError(f"Texture group {group.description!r} contains no usable images.")

        output_mode = _common_mode(list(opened.values()))
        output = Image.new(output_mode, expected_size, 0)
        try:
            missing: list[str] = []
            for material in materials:
                key = material.name.casefold()
                source = opened.get(key)
                if source is None:
                    missing.append(material.name)
                    continue
                converted = None
                if source.mode != output_mode:
                    converted = source.convert(output_mode)
                    source = converted
                try:
                    output.paste(source, (0, 0), masks[key])
                finally:
                    if converted is not None:
                        converted.close()

            output_path.parent.mkdir(parents=True, exist_ok=True)
            output.save(output_path)
            return missing
        finally:
            output.close()
    finally:
        for image in opened.values():
            image.close()


def combine_all(
    id_map_path: Path,
    legend_path: Path,
    input_dir: Path,
    output_dir: Path,
    aliases: dict[str, str] | None = None,
    recursive: bool = False,
    tolerance: int | None = None,
    overwrite: bool = False,
    dry_run: bool = False,
) -> tuple[list[Path], list[Path]]:
    _require_pillow()
    if not input_dir.is_dir():
        raise CombineError(f"Input directory does not exist: {input_dir}")
    if tolerance is not None and (tolerance < 0 or tolerance > 32):
        raise CombineError("Color tolerance must be between 0 and 32.")

    materials = load_legend(legend_path, aliases)
    groups, unmatched = discover_groups(
        input_dir,
        materials,
        recursive=recursive,
        excluded_paths=(id_map_path,),
    )
    if not groups:
        names = ", ".join(material.filename_name for material in materials)
        raise CombineError(
            f"No supported textures in {input_dir} contained a legend material name "
            f"({names}). Use --material-alias when Substance texture-set names differ."
        )

    outputs: list[Path] = []
    if dry_run:
        for group in groups:
            outputs.append(output_dir / group.output_parent / group.output_name)
        return outputs, unmatched

    try:
        with Image.open(id_map_path) as loaded_id_map:
            id_map = loaded_id_map.convert("RGB")
            id_map.load()
    except OSError as exc:
        raise CombineError(f"Could not load ID map {id_map_path}: {exc}") from exc

    try:
        masks = {}
        for material in materials:
            material_tolerance = tolerance if tolerance is not None else 0
            mask = _mask_for_color(id_map, material.rgb, material_tolerance)
            if tolerance is None and mask.getbbox() is None:
                # Blender stores byte images after float-to-byte conversion.
                # Values exactly between bytes can differ from ordinary Python
                # rounding by one, so only widen a mask when exact pixels do
                # not exist at all.
                material_tolerance = 1
                mask = _mask_for_color(id_map, material.rgb, material_tolerance)
                if mask.getbbox() is not None:
                    print(
                        f"notice: matched {material.name!r} with automatic ±1 byte tolerance",
                        file=sys.stderr,
                    )
            masks[material.name.casefold()] = mask
        occupied = Image.new("1", id_map.size, 0)
        for material in materials:
            mask = masks[material.name.casefold()]
            if ImageChops.multiply(occupied, mask).getbbox() is not None:
                raise CombineError(
                    f"The mask for {material.name!r} overlaps another material mask. "
                    "Lower --tolerance (exact matching is --tolerance 0)."
                )
            occupied = ImageChops.lighter(occupied, mask)
        empty_masks = [
            material.name
            for material in materials
            if masks[material.name.casefold()].getbbox() is None
        ]
        if empty_masks:
            print(
                "warning: no ID-map pixels matched: " + ", ".join(empty_masks),
                file=sys.stderr,
            )

        output_paths = [
            output_dir / group.output_parent / group.output_name
            for group in groups
        ]
        if len(set(output_paths)) != len(output_paths):
            duplicates = sorted(
                {path for path in output_paths if output_paths.count(path) > 1},
                key=str,
            )
            raise CombineError(
                "Multiple discovered groups resolve to the same output: "
                + ", ".join(str(path) for path in duplicates)
            )

        for group, output_path in zip(groups, output_paths, strict=True):
            missing = combine_group(
                group,
                materials,
                masks,
                id_map.size,
                output_path,
                overwrite=overwrite,
            )
            if missing:
                print(
                    f"warning: {group.description} has no texture for: {', '.join(missing)}",
                    file=sys.stderr,
                )
            outputs.append(output_path)
    finally:
        id_map.close()

    return outputs, unmatched


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Combine Substance Painter textures exported as separate material/texture "
            "sets, using a Material ID bake and JSON legend as exact masks."
        )
    )
    parser.add_argument("--id-map", required=True, type=Path, help="Material ID PNG")
    parser.add_argument(
        "--legend",
        required=True,
        type=Path,
        help="JSON legend written beside the Material ID PNG",
    )
    parser.add_argument(
        "--input-dir",
        required=True,
        type=Path,
        help="Directory containing Substance Painter exports",
    )
    parser.add_argument(
        "--output-dir",
        required=True,
        type=Path,
        help="Directory for combined channel textures",
    )
    parser.add_argument(
        "--material-alias",
        action="append",
        default=[],
        metavar="LEGEND=FILENAME",
        help="Map a Blender material name to a different Substance filename token",
    )
    parser.add_argument("--recursive", action="store_true", help="Search input subdirectories")
    parser.add_argument(
        "--tolerance",
        type=int,
        default=None,
        help=(
            "Fixed per-channel ID color tolerance in 8-bit values; by default the tool "
            "tries exact matching and a ±1 fallback only when exact pixels are absent"
        ),
    )
    parser.add_argument("--overwrite", action="store_true", help="Replace existing outputs")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print discovered groups and output paths without reading or writing images",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        aliases = _parse_aliases(args.material_alias)
        outputs, unmatched = combine_all(
            id_map_path=args.id_map.resolve(),
            legend_path=args.legend.resolve(),
            input_dir=args.input_dir.resolve(),
            output_dir=args.output_dir.resolve(),
            aliases=aliases,
            recursive=args.recursive,
            tolerance=args.tolerance,
            overwrite=args.overwrite,
            dry_run=args.dry_run,
        )
    except CombineError as exc:
        parser.exit(2, f"error: {exc}\n")

    action = "Would write" if args.dry_run else "Wrote"
    for output in outputs:
        print(f"{action}: {output}")
    if unmatched:
        print(f"Skipped {len(unmatched)} image(s) with no material-name match:")
        for path in unmatched:
            print(f"  {path}")
    print(f"{len(outputs)} texture channel(s) {'discovered' if args.dry_run else 'combined'}.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())