#!/usr/bin/env python3
"""Build the bounded FG-3 Siltwake runtime art package from retained media."""

from __future__ import annotations

import argparse
from collections import deque
import hashlib
import json
from pathlib import Path

from PIL import Image


ROOT = Path(__file__).resolve().parents[1]
MEDIA_ROOT = ROOT / "docs/runtime/fg3/media"
RUNTIME_ROOT = ROOT / "runtime/assets/fg3"
ENVIRONMENT_SOURCE = MEDIA_ROOT / "raw/siltwake-layers.png"
BOSS_ACTION_SOURCE = MEDIA_ROOT / "raw/rg02-boss-actions.png"
IDENTITY_SOURCE = MEDIA_ROOT / "raw/rg02-company-enemy-identities.png"
TARGET_SIZES = ((380, 180), (422, 195), (320, 180))
LAYER_NAMES = ("background", "midground", "foreground")
PALETTE_SIZE = 32
IDENTITY_LOGICAL_SIZE = 36
IDENTITY_CELL_SIZE = 72
IDENTITY_PALETTE_SIZE = 16
BOSSES = (
    {
        "identity_id": "BO-R2-01",
        "slug": "sluiceback_mother",
        "source_row": 0,
        "state": "boss_action",
        "fps": 8,
        "loop": False,
        "play_size": [78, 58],
        "pivot": [39, 56],
        "cell_size": [128, 96],
        "provider_task_id": "task_01M056P0CCPYVFS2GXD91GK7KZ",
    },
    {
        "identity_id": "BO-R2-02",
        "slug": "drowned_bell",
        "source_row": 1,
        "state": "boss_action",
        "fps": 8,
        "loop": False,
        "play_size": [62, 72],
        "pivot": [31, 70],
        "cell_size": [96, 96],
        "provider_task_id": "task_01M056P0CCPYVFS2GXD91GK7KZ",
    },
)

IDENTITIES = (
    {"identity_id": "PF-B09", "slug": "coilsmith", "source_cell": 0, "play_size": [36, 40]},
    {"identity_id": "PF-B07", "slug": "kindhand", "source_cell": 1, "play_size": [34, 40]},
    {"identity_id": "EN-R2-01", "slug": "reed_hook", "source_cell": 2, "play_size": [34, 40]},
    {"identity_id": "EN-R2-02", "slug": "mudlamp_leech", "source_cell": 3, "play_size": [38, 32]},
    {"identity_id": "EN-R2-03", "slug": "silt_mimic", "source_cell": 4, "play_size": [36, 34]},
    {"identity_id": "EN-R2-04", "slug": "bank_tortoise", "source_cell": 5, "play_size": [38, 32]},
    {"identity_id": "EN-R2-05", "slug": "fen_piper", "source_cell": 6, "play_size": [34, 40]},
    {"identity_id": "EN-R2-06", "slug": "brackwater_wisp", "source_cell": 7, "play_size": [32, 38]},
    {"identity_id": "EN-R2-07", "slug": "ferry_crab", "source_cell": 8, "play_size": [40, 32]},
    {"identity_id": "EN-R2-08", "slug": "drownseed_pod", "source_cell": 9, "play_size": [34, 36]},
    {"identity_id": "EN-R2-09", "slug": "mire_medic", "source_cell": 10, "play_size": [34, 40]},
    {"identity_id": "EN-R2-10", "slug": "fogjaw", "source_cell": 11, "play_size": [40, 32]},
    {"identity_id": "EN-R2-11", "slug": "channel_widow", "source_cell": 12, "play_size": [36, 42]},
    {"identity_id": "EN-R2-12", "slug": "sluice_elder", "source_cell": 13, "play_size": [36, 42]},
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--rows-only",
        action="store_true",
        help="write the two deterministic raw action rows, then stop",
    )
    return parser.parse_args()


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def remove_magenta(image: Image.Image) -> Image.Image:
    rgba = image.convert("RGBA")
    cleaned = []
    for red, green, blue, alpha in rgba.getdata():
        magenta = (
            red >= 140
            and blue >= 140
            and green <= 165
            and min(red, blue) - green >= 45
            and abs(red - blue) <= 110
        )
        cleaned.append((red, green, blue, 0 if magenta else alpha))
    rgba.putdata(cleaned)
    return rgba


def crop_visible(image: Image.Image) -> Image.Image:
    bounds = image.getchannel("A").getbbox()
    if bounds is None:
        raise ValueError("FG-3 source layer has no non-key pixels")
    return image.crop(bounds)


def resized_cover(image: Image.Image, size: tuple[int, int]) -> Image.Image:
    scale = max(size[0] / image.width, size[1] / image.height)
    resized = image.resize(
        (max(1, round(image.width * scale)), max(1, round(image.height * scale))),
        Image.Resampling.BOX,
    )
    left = max(0, (resized.width - size[0]) // 2)
    top = max(0, (resized.height - size[1]) // 2)
    return resized.crop((left, top, left + size[0], top + size[1]))


def resized_contain(image: Image.Image, size: tuple[int, int]) -> Image.Image:
    scale = min(size[0] / image.width, size[1] / image.height)
    return image.resize(
        (max(1, round(image.width * scale)), max(1, round(image.height * scale))),
        Image.Resampling.BOX,
    )


def build_base_layers(size: tuple[int, int], bands: list[Image.Image]) -> dict[str, Image.Image]:
    base_size = (size[0] // 2, (size[1] + 1) // 2)
    background = Image.new("RGBA", base_size, (7, 26, 44, 255))
    background.alpha_composite(resized_cover(bands[0], base_size))

    midground = Image.new("RGBA", base_size, (0, 0, 0, 0))
    middle = resized_contain(bands[1], base_size)
    midground.alpha_composite(
        middle,
        ((base_size[0] - middle.width) // 2, base_size[1] - middle.height),
    )

    foreground = Image.new("RGBA", base_size, (0, 0, 0, 0))
    front = resized_contain(bands[2], base_size)
    foreground.alpha_composite(
        front,
        ((base_size[0] - front.width) // 2, base_size[1] - front.height),
    )
    return {"background": background, "midground": midground, "foreground": foreground}


def make_shared_palette(layers: dict[str, Image.Image]) -> Image.Image:
    samples: list[tuple[int, int, int]] = []
    for layer in layers.values():
        samples.extend((r, g, b) for r, g, b, a in layer.getdata() if a > 0)
    stride = max(1, len(samples) // 100000)
    selected = samples[::stride]
    source = Image.new("RGB", (len(selected), 1))
    source.putdata(selected)
    return source.quantize(
        colors=PALETTE_SIZE,
        method=Image.Quantize.MEDIANCUT,
        dither=Image.Dither.NONE,
    )


def apply_palette(image: Image.Image, palette: Image.Image) -> Image.Image:
    alpha = image.getchannel("A").point(lambda value: 255 if value >= 128 else 0)
    limited = image.convert("RGB").quantize(
        palette=palette,
        dither=Image.Dither.NONE,
    ).convert("RGBA")
    limited.putalpha(alpha)
    return limited


def unique_colors(image: Image.Image) -> int:
    return len({pixel for pixel in image.getdata() if pixel[3] > 0})


def remove_tiny_components(image: Image.Image, min_area: int = 4) -> Image.Image:
    rgba = image.convert("RGBA")
    alpha = list(rgba.getchannel("A").getdata())
    width, height = rgba.size
    visited = bytearray(width * height)
    retained = bytearray(alpha)
    for start, value in enumerate(alpha):
        if value == 0 or visited[start]:
            continue
        visited[start] = 1
        queue: deque[int] = deque([start])
        component: list[int] = []
        while queue:
            current = queue.popleft()
            component.append(current)
            x, y = current % width, current // width
            for next_y in range(max(0, y - 1), min(height, y + 2)):
                row_start = next_y * width
                for next_x in range(max(0, x - 1), min(width, x + 2)):
                    neighbor = row_start + next_x
                    if alpha[neighbor] and not visited[neighbor]:
                        visited[neighbor] = 1
                        queue.append(neighbor)
        if len(component) < min_area:
            for pixel in component:
                retained[pixel] = 0
    cleaned_alpha = Image.new("L", rgba.size)
    cleaned_alpha.putdata(retained)
    rgba.putalpha(cleaned_alpha)
    return rgba


def fit_identity(image: Image.Image) -> Image.Image:
    visible = crop_visible(remove_magenta(image))
    content_limit = IDENTITY_LOGICAL_SIZE - 6
    scale = min(content_limit / visible.width, content_limit / visible.height)
    resized = visible.resize(
        (max(1, round(visible.width * scale)), max(1, round(visible.height * scale))),
        Image.Resampling.BOX,
    )
    alpha = resized.getchannel("A").point(lambda value: 255 if value >= 128 else 0)
    limited = resized.convert("RGB").quantize(
        colors=IDENTITY_PALETTE_SIZE,
        method=Image.Quantize.MEDIANCUT,
        dither=Image.Dither.NONE,
    ).convert("RGBA")
    limited.putalpha(alpha)
    logical = Image.new(
        "RGBA",
        (IDENTITY_LOGICAL_SIZE, IDENTITY_LOGICAL_SIZE),
        (0, 0, 0, 0),
    )
    left = (IDENTITY_LOGICAL_SIZE - limited.width) // 2
    top = IDENTITY_LOGICAL_SIZE - 3 - limited.height
    logical.alpha_composite(limited, (left, top))
    logical = remove_tiny_components(logical)
    return logical.resize(
        (IDENTITY_CELL_SIZE, IDENTITY_CELL_SIZE),
        Image.Resampling.NEAREST,
    )


def prepare_boss_rows() -> list[dict[str, object]]:
    source = Image.open(BOSS_ACTION_SOURCE).convert("RGBA")
    records: list[dict[str, object]] = []
    for boss in BOSSES:
        row_index = int(boss["source_row"])
        top = round(row_index * source.height / len(BOSSES))
        bottom = round((row_index + 1) * source.height / len(BOSSES))
        run_dir = MEDIA_ROOT / "runs" / str(boss["slug"])
        output = run_dir / "raw" / f"{boss['state']}.png"
        output.parent.mkdir(parents=True, exist_ok=True)
        source.crop((0, top, source.width, bottom)).save(
            output,
            format="PNG",
            optimize=False,
            compress_level=9,
        )
        records.append(
            {
                "identity_id": boss["identity_id"],
                "source_row": row_index,
                "source_rect": [0, top, source.width, bottom - top],
                "path": output.relative_to(ROOT).as_posix(),
                "sha256": sha256(output),
            }
        )
    rows_manifest = MEDIA_ROOT / "runs/boss-action-rows.json"
    rows_manifest.write_text(
        json.dumps(
            {
                "kind": "abg-fg3-boss-action-rows",
                "source": BOSS_ACTION_SOURCE.relative_to(ROOT).as_posix(),
                "source_sha256": sha256(BOSS_ACTION_SOURCE),
                "source_dimensions": list(source.size),
                "rows": records,
            },
            ensure_ascii=True,
            indent=2,
            sort_keys=True,
        )
        + "\n",
        encoding="utf-8",
    )
    return records


def process_environment() -> dict[str, object]:
    source = Image.open(ENVIRONMENT_SOURCE).convert("RGBA")
    cuts = (0, source.height // 3, (source.height * 2) // 3, source.height)
    bands = [
        crop_visible(
            remove_magenta(
                source.crop((0, cuts[index], source.width, cuts[index + 1]))
            )
        )
        for index in range(3)
    ]
    reference_layers = build_base_layers((422, 195), bands)
    palette = make_shared_palette(reference_layers)
    records: dict[str, dict[str, object]] = {name: {} for name in LAYER_NAMES}

    for width, height in TARGET_SIZES:
        size_key = f"{width}x{height}"
        base_layers = build_base_layers((width, height), bands)
        output_dir = RUNTIME_ROOT / "environment" / size_key
        output_dir.mkdir(parents=True, exist_ok=True)
        for name in LAYER_NAMES:
            limited = apply_palette(base_layers[name], palette)
            rendered = limited.resize(
                (width, limited.height * 2),
                Image.Resampling.NEAREST,
            ).crop((0, 0, width, height))
            output = output_dir / f"{name}.png"
            rendered.save(output, format="PNG", optimize=False, compress_level=9)
            records[name][size_key] = {
                "path": output.relative_to(ROOT / "runtime").as_posix(),
                "dimensions": [width, height],
                "logical_source_dimensions": list(limited.size),
                "cluster_pitch": 2,
                "colors": unique_colors(rendered),
                "sha256": sha256(output),
            }

    return {
        "asset_id": "FG3-BG-SILTWAKE-01",
        "region_id": "RG-02",
        "source": ENVIRONMENT_SOURCE.relative_to(ROOT).as_posix(),
        "source_sha256": sha256(ENVIRONMENT_SOURCE),
        "source_dimensions": list(source.size),
        "source_band_cuts": list(cuts),
        "provider_task_id": "task_01M056E56CNT64A0M3VB617Z9R",
        "palette_size": PALETTE_SIZE,
        "layers": records,
    }


def process_bosses() -> dict[str, dict[str, object]]:
    records: dict[str, dict[str, object]] = {}
    for boss in BOSSES:
        run_dir = MEDIA_ROOT / "runs" / str(boss["slug"])
        manifest_path = run_dir / "manifest.json"
        if not manifest_path.is_file():
            raise FileNotFoundError(f"sprite-gen manifest missing: {manifest_path}")
        source_manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        state = str(boss["state"])
        frame_rects = source_manifest["frame_layout"]["rows"][state]
        cell_width, cell_height = (int(value) for value in boss["cell_size"])
        if len(frame_rects) != 4:
            raise ValueError(f"{boss['identity_id']} must have four action frames")

        output_dir = RUNTIME_ROOT / "bosses" / str(boss["slug"])
        frame_dir = output_dir / state
        frame_dir.mkdir(parents=True, exist_ok=True)
        rendered_frames: list[Image.Image] = []
        runtime_frames: list[dict[str, object]] = []
        for index, _frame_rect in enumerate(frame_rects):
            source_frame = run_dir / "frames" / state / f"frame-{index}.png"
            with Image.open(source_frame) as opened:
                frame = opened.convert("RGBA")
            if frame.size != (cell_width, cell_height):
                raise ValueError(f"unexpected Boss frame size: {source_frame} {frame.size}")
            alpha = frame.getchannel("A").point(lambda value: 255 if value >= 128 else 0)
            frame.putalpha(alpha)
            output_frame = frame_dir / f"frame-{index}.png"
            frame.save(output_frame, format="PNG", optimize=False, compress_level=9)
            rendered_frames.append(frame)
            runtime_frames.append(
                {
                    "path": output_frame.relative_to(ROOT / "runtime").as_posix(),
                    "dimensions": list(frame.size),
                    "atlas_rect": [index * cell_width, 0, cell_width, cell_height],
                    "sha256": sha256(output_frame),
                }
            )

        atlas_output = output_dir / "atlas.png"
        atlas = Image.new(
            "RGBA",
            (cell_width * len(rendered_frames), cell_height),
            (0, 0, 0, 0),
        )
        for index, frame in enumerate(rendered_frames):
            atlas.alpha_composite(frame, (index * cell_width, 0))
        atlas.save(atlas_output, format="PNG", optimize=False, compress_level=9)

        raw_row = run_dir / "raw" / f"{state}.png"
        records[str(boss["identity_id"])] = {
            "asset_id": f"FG3-BOSS-{str(boss['identity_id']).replace('-', '')}-01",
            "source": BOSS_ACTION_SOURCE.relative_to(ROOT).as_posix(),
            "source_sha256": sha256(BOSS_ACTION_SOURCE),
            "source_row": int(boss["source_row"]),
            "raw_row": raw_row.relative_to(ROOT).as_posix(),
            "raw_row_sha256": sha256(raw_row),
            "sprite_manifest": manifest_path.relative_to(ROOT).as_posix(),
            "sprite_manifest_sha256": sha256(manifest_path),
            "provider_task_id": boss["provider_task_id"],
            "state": state,
            "fps": int(boss["fps"]),
            "loop": bool(boss["loop"]),
            "pivot": list(boss["pivot"]),
            "play_size": list(boss["play_size"]),
            "atlas": atlas_output.relative_to(ROOT / "runtime").as_posix(),
            "atlas_sha256": sha256(atlas_output),
            "frames": runtime_frames,
        }
    return records


def process_identities() -> dict[str, dict[str, object]]:
    source = Image.open(IDENTITY_SOURCE).convert("RGBA")
    records: dict[str, dict[str, object]] = {}
    for identity in IDENTITIES:
        source_cell = int(identity["source_cell"])
        row, column = divmod(source_cell, 4)
        bounds = (
            round(column * source.width / 4),
            round(row * source.height / 4),
            round((column + 1) * source.width / 4),
            round((row + 1) * source.height / 4),
        )
        frame = fit_identity(source.crop(bounds))
        output_dir = RUNTIME_ROOT / "actors" / str(identity["slug"])
        frame_dir = output_dir / "idle"
        frame_dir.mkdir(parents=True, exist_ok=True)
        output_frame = frame_dir / "frame-0.png"
        atlas_output = output_dir / "atlas.png"
        frame.save(output_frame, format="PNG", optimize=False, compress_level=9)
        frame.save(atlas_output, format="PNG", optimize=False, compress_level=9)
        identity_id = str(identity["identity_id"])
        records[identity_id] = {
            "asset_id": f"FG3-ACTOR-{identity_id.replace('-', '')}-01",
            "source": IDENTITY_SOURCE.relative_to(ROOT).as_posix(),
            "source_sha256": sha256(IDENTITY_SOURCE),
            "source_cell": source_cell,
            "source_rect": [bounds[0], bounds[1], bounds[2] - bounds[0], bounds[3] - bounds[1]],
            "provider_task_id": "task_01M058059WFB3F4G96SC9SDRWD",
            "state": "idle",
            "fps": 1,
            "loop": True,
            "pivot": [int(identity["play_size"][0]) // 2, int(identity["play_size"][1]) - 2],
            "play_size": list(identity["play_size"]),
            "logical_grid": [IDENTITY_LOGICAL_SIZE, IDENTITY_LOGICAL_SIZE],
            "palette_size": IDENTITY_PALETTE_SIZE,
            "colors": unique_colors(frame),
            "atlas": atlas_output.relative_to(ROOT / "runtime").as_posix(),
            "atlas_sha256": sha256(atlas_output),
            "frames": [
                {
                    "path": output_frame.relative_to(ROOT / "runtime").as_posix(),
                    "dimensions": list(frame.size),
                    "atlas_rect": [0, 0, IDENTITY_CELL_SIZE, IDENTITY_CELL_SIZE],
                    "sha256": sha256(output_frame),
                }
            ],
        }
    return records


def main() -> int:
    args = parse_args()
    rows = prepare_boss_rows()
    if args.rows_only:
        print(
            json.dumps(
                {"ok": True, "rows": len(rows), "mode": "rows-only"},
                sort_keys=True,
            )
        )
        return 0

    RUNTIME_ROOT.mkdir(parents=True, exist_ok=True)
    manifest = {
        "version": 1,
        "kind": "abg-fg3-runtime-media",
        "license": "owner-authorized original generation for Aetherbound Guild",
        "filter": "nearest",
        "integer_placement": True,
        "environment": process_environment(),
        "actors": process_identities(),
        "bosses": process_bosses(),
    }
    manifest_path = RUNTIME_ROOT / "asset-manifest.json"
    manifest_path.write_text(
        json.dumps(manifest, ensure_ascii=True, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(
        json.dumps(
            {
                "ok": True,
                "manifest": str(manifest_path),
                "bosses": len(manifest["bosses"]),
                "actors": len(manifest["actors"]),
                "environment_sizes": len(TARGET_SIZES),
            },
            sort_keys=True,
        )
    )
    return 0


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