# python pox_tool.py -export   "C:\Program Files (x86)\Steam\steamapps\common\Siege of Avalon Anthology\ArtLib\Resources\StaticObject\Containers\Barrels"
# python pox_tool.py -generate "C:\Program Files (x86)\Steam\steamapps\common\Siege of Avalon Anthology\ArtLib\Resources\StaticObject\Containers\Barrels" --alpha-threshold 255


# python pox_tool.py -generate "C:\Program Files (x86)\Steam\steamapps\common\Siege of Avalon Anthology\ArtLib\Resources\StaticObject\Containers\Barrels" --alpha-mode solidify_bleed --alpha-drop 10 --bleed-radius 3



#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
POX Tool (Siege of Avalon)
==========================

EXPORT (recursive):
- *.pox -> <pox_basename>/<basename>.ini + <basename>_001.png...

GENERATE (recursive):
- <folder with 1 ini + png frames> -> writes .pox one level above, then deletes folder.

Additions:
- alpha handling modes for GENERATE to avoid white/bright halos in formats without proper alpha:
    --alpha-mode threshold_delete   : any alpha < alpha-drop => 0, else 255 (hard edges / shrink possible)
    --alpha-mode solidify          : any alpha > 0 becomes 255 (keeps silhouette, may keep halos if RGB is bad)
    --alpha-mode solidify_bleed    : semi-transparent (0<a<255) pixels get RGB from nearby opaque pixels, then a=255 (recommended)
- --alpha-drop N : threshold for treating small alpha as 0 (default 10 for solidify modes, 255 for threshold_delete is typical)
- --bleed-radius R : neighbor radius for RGB reconstruction (default 2)
"""

import argparse
import os
import re
import shutil
import struct
from dataclasses import dataclass
from typing import List, Optional, Tuple

from PIL import Image


# ----------------- Shared helpers --------------------------------------------

_ns_re = re.compile(r"(\d+)")

def natural_key(path: str):
    base = os.path.basename(path)
    parts = _ns_re.split(base)
    out = []
    for p in parts:
        out.append(int(p) if p.isdigit() else p.lower())
    return out


def ensure_dir(p: str) -> None:
    os.makedirs(p, exist_ok=True)


def read_text_cp1252(path: str) -> str:
    with open(path, "r", encoding="cp1252", errors="replace") as f:
        return f.read()


def write_text_cp1252(path: str, text: str) -> None:
    with open(path, "w", encoding="cp1252", errors="replace") as f:
        f.write(text)


# ----------------- POX structures (matching POXStudio) -----------------------

RLEHDR_STRUCT = struct.Struct("<iiIIiiIi")  # 32 bytes


@dataclass
class RLEHDR:
    srcx: int
    srcy: int
    wdh: int
    hgh: int
    adjx: int
    adjy: int
    pixfmt: int
    dataptr: int

    def pack(self) -> bytes:
        return RLEHDR_STRUCT.pack(
            self.srcx, self.srcy, self.wdh, self.hgh,
            self.adjx, self.adjy, self.pixfmt, self.dataptr
        )

    @staticmethod
    def unpack(buf: bytes, offset: int) -> "RLEHDR":
        return RLEHDR(*RLEHDR_STRUCT.unpack_from(buf, offset))


# ----------------- Color conversion ------------------------------------------

def rgb565_to_rgb888(v: int) -> Tuple[int, int, int]:
    # POXStudio style: shift up (not bit replication)
    r = ((v & 0xF800) >> 11) << 3
    g = ((v & 0x07E0) >> 5) << 2
    b = (v & 0x001F) << 3
    return (r & 0xFF, g & 0xFF, b & 0xFF)


def color_to_rgb565(r: int, g: int, b: int) -> int:
    rr = (r & 0xFF) >> 3
    gg = (g & 0xFF) >> 2
    bb = (b & 0xFF) >> 3
    return (rr << 11) | (gg << 5) | bb


# ----------------- INI helpers -----------------------------------------------

def ini_get_int(ini_text: str, key: str) -> Optional[int]:
    m = re.search(rf"(?im)^\s*{re.escape(key)}\s*=\s*(-?\d+)\s*$", ini_text)
    return int(m.group(1)) if m else None


def parse_transparent_color_from_ini(ini_text: str) -> Optional[Tuple[int, int, int]]:
    m = re.search(r"(?im)^\s*TransparentColor\s*=\s*(\d+)\s*$", ini_text)
    if not m:
        return None
    dec = int(m.group(1))
    r = (dec >> 16) & 0xFF
    g = (dec >> 8) & 0xFF
    b = dec & 0xFF
    return (r, g, b)


def extract_type_marker(ini_text: str) -> Optional[str]:
    m = re.search(r"(?im)^\s*;\s*POX_TOOL_TYPE\s*=\s*([A-Za-z0-9]{2})\s*$", ini_text)
    return m.group(1).upper() if m else None


def inject_type_marker(ini_text: str, res_type: str) -> str:
    res_type = res_type.upper()
    marker = f";POX_TOOL_TYPE={res_type}\n"
    if re.search(r"(?im)^\s*;\s*POX_TOOL_TYPE\s*=\s*[A-Za-z0-9]{2}\s*$", ini_text):
        return ini_text
    return marker + ini_text


# ----------------- Export (POX -> PNG/INI) -----------------------------------

def derive_canvas_size_from_headers(headers: List[RLEHDR]) -> Tuple[int, int]:
    max_w = 1
    max_h = 1
    for h in headers:
        max_w = max(max_w, h.adjx + int(h.wdh))
        max_h = max(max_h, h.adjy + int(h.hgh))
    return max_w, max_h


def read_pox(path: str) -> Tuple[str, str, List[RLEHDR], bytes]:
    data = open(path, "rb").read()

    if len(data) < 12:
        raise ValueError("File too small to be a POX.")

    if data[:4] != b"POXA":
        raise ValueError("Not a POXA file (header mismatch).")

    res_type = data[4:6].decode("ascii", errors="replace")
    ini_len = struct.unpack_from("<I", data, 8)[0]

    ini_start = 12
    ini_end = ini_start + ini_len
    if ini_end > len(data):
        raise ValueError("INI length exceeds file size (corrupt POX).")

    ini_text = data[ini_start:ini_end].decode("cp1252", errors="replace")

    off = ini_end
    if off + 2 > len(data):
        raise ValueError("Unexpected EOF after INI.")

    (bb1,) = struct.unpack_from("<H", data, off)
    off += 2
    if bb1 != 0x4242:
        raise ValueError("Missing BB marker after INI (unexpected format).")

    pic_cnt = struct.unpack_from("<I", data, off)[0]
    off += 4
    rle_size = struct.unpack_from("<I", data, off)[0]
    off += 4

    need_hdr_bytes = pic_cnt * RLEHDR_STRUCT.size
    if off + need_hdr_bytes > len(data):
        raise ValueError("Not enough bytes for all RLE headers (corrupt POX).")

    headers = [RLEHDR.unpack(data, off + i * RLEHDR_STRUCT.size) for i in range(pic_cnt)]
    off += need_hdr_bytes

    if off + rle_size > len(data):
        raise ValueError("RLE data exceeds file size (corrupt POX).")

    rle_data = data[off:off + rle_size]
    return res_type, ini_text, headers, rle_data


def decode_frame_rgba(
    rle_data: bytes,
    start_ptr: int,
    hdr: RLEHDR,
    canvas_w: int,
    canvas_h: int,
    key_rgb: Optional[Tuple[int, int, int]]
) -> Image.Image:
    if hdr.pixfmt != 2:
        raise ValueError(f"Unsupported PixFmt={hdr.pixfmt} (only 2/RGB565 supported)")

    img = Image.new("RGBA", (canvas_w, canvas_h), (0, 0, 0, 0))
    px = img.load()

    data_len = len(rle_data)
    p = start_ptr

    x = 0
    y = 0

    max_steps = data_len + 1024
    steps = 0

    while 0 <= p < data_len and steps < max_steps:
        steps += 1
        cmd = rle_data[p]
        p += 1

        if cmd == 0x00 or cmd >= 0x04:
            break

        if cmd == 0x03:
            y += 1
            continue

        if cmd == 0x02:
            if p + 4 > data_len:
                break
            (i,) = struct.unpack_from("<i", rle_data, p)
            p += 4
            x += (i // 2)
            continue

        if cmd == 0x01:
            if p + 4 > data_len:
                break
            (count,) = struct.unpack_from("<i", rle_data, p)
            p += 4
            if count <= 0:
                continue

            need = count * 2
            if p + need > data_len:
                break

            for _ in range(count):
                (c565,) = struct.unpack_from("<H", rle_data, p)
                p += 2
                r, g, b = rgb565_to_rgb888(c565)

                if key_rgb is None or (r, g, b) != key_rgb:
                    xx = x + hdr.adjx
                    yy = y + hdr.adjy
                    if 0 <= xx < canvas_w and 0 <= yy < canvas_h:
                        px[xx, yy] = (r, g, b, 255)

                x += 1

    return img


def export_one_pox(pox_path: str) -> None:
    res_type, ini_text, headers, rle_data = read_pox(pox_path)
    key_rgb = parse_transparent_color_from_ini(ini_text)

    p = os.path.abspath(pox_path)
    pdir = os.path.dirname(p)
    base = os.path.splitext(os.path.basename(p))[0]

    out_dir = os.path.join(pdir, base)
    ensure_dir(out_dir)

    ini_text_marked = inject_type_marker(ini_text, res_type)
    ini_out = os.path.join(out_dir, f"{base}.ini")
    write_text_cp1252(ini_out, ini_text_marked)

    w = ini_get_int(ini_text, "ImageWidth")
    h = ini_get_int(ini_text, "ImageHeight")
    if w is None or h is None or w <= 0 or h <= 0:
        w, h = derive_canvas_size_from_headers(headers)

    if len(headers) > 999:
        raise ValueError(f"{pox_path}: frame count {len(headers)} exceeds 999 (requested limit).")

    base_ptr = headers[0].dataptr

    for i, hdr in enumerate(headers, start=1):
        start_ptr_norm = hdr.dataptr - base_ptr
        if not (0 <= start_ptr_norm < len(rle_data)):
            start_ptr = hdr.dataptr
        else:
            start_ptr = start_ptr_norm

        img = decode_frame_rgba(rle_data, start_ptr, hdr, w, h, key_rgb)
        png_name = f"{base}_{i:03}.png"
        img.save(os.path.join(out_dir, png_name), format="PNG")


def export_recursive(root: str) -> None:
    root = os.path.abspath(root)
    for dirpath, _, filenames in os.walk(root):
        for fn in filenames:
            if fn.lower().endswith(".pox"):
                pox_path = os.path.join(dirpath, fn)
                try:
                    export_one_pox(pox_path)
                    print(f"EXPORT OK: {pox_path}")
                except Exception as e:
                    print(f"EXPORT FAIL: {pox_path} :: {e}")


# ----------------- Generate (PNG/INI -> POX) ---------------------------------

def apply_keycolor_transparency(img_rgba: Image.Image, key_rgb: Optional[Tuple[int, int, int]]) -> Image.Image:
    if img_rgba.mode != "RGBA":
        img_rgba = img_rgba.convert("RGBA")
    if key_rgb is None:
        return img_rgba

    px = img_rgba.load()
    w, h = img_rgba.size
    kr, kg, kb = key_rgb
    for y in range(h):
        for x in range(w):
            r, g, b, a = px[x, y]
            if (r, g, b) == (kr, kg, kb):
                px[x, y] = (r, g, b, 0)
    return img_rgba


def alpha_mode_threshold_delete(img: Image.Image, alpha_drop: int) -> Image.Image:
    """
    Hard threshold:
      - a < alpha_drop => a=0
      - else => a=255
    """
    if img.mode != "RGBA":
        img = img.convert("RGBA")
    px = img.load()
    w, h = img.size
    for y in range(h):
        for x in range(w):
            r, g, b, a = px[x, y]
            if a < alpha_drop:
                px[x, y] = (r, g, b, 0)
            else:
                px[x, y] = (r, g, b, 255)
    return img


def alpha_mode_solidify(img: Image.Image, alpha_drop: int) -> Image.Image:
    """
    Keeps silhouette (no shrinking):
      - a < alpha_drop => a=0
      - else if a>0 => a=255 (RGB unchanged)
    This can still keep halos if RGB on semi-transparent pixels is bad.
    """
    if img.mode != "RGBA":
        img = img.convert("RGBA")
    px = img.load()
    w, h = img.size
    for y in range(h):
        for x in range(w):
            r, g, b, a = px[x, y]
            if a < alpha_drop:
                px[x, y] = (r, g, b, 0)
            elif a != 0:
                px[x, y] = (r, g, b, 255)
    return img


def alpha_mode_solidify_bleed(img: Image.Image, alpha_drop: int, radius: int) -> Image.Image:
    """
    Recommended for formats without proper alpha:
      - a < alpha_drop => transparent (0)
      - 0 < a < 255 => replace RGB using nearby opaque pixels (a==255), then set a=255
      - a==255 stays
    This removes bright/white fringes while keeping silhouette (no shrink).
    """
    if img.mode != "RGBA":
        img = img.convert("RGBA")

    px = img.load()
    w, h = img.size

    # Snapshot alpha so neighbor decisions don't change mid-loop
    alpha = [[px[x, y][3] for x in range(w)] for y in range(h)]

    for y in range(h):
        for x in range(w):
            a = alpha[y][x]
            r, g, b, _ = px[x, y]

            if a < alpha_drop:
                px[x, y] = (r, g, b, 0)
                continue

            if a == 0:
                # keep transparent
                continue

            if a == 255:
                continue

            # 0 < a < 255 : reconstruct RGB from nearby opaque pixels
            samples = []
            for dy in range(-radius, radius + 1):
                ny = y + dy
                if ny < 0 or ny >= h:
                    continue
                for dx in range(-radius, radius + 1):
                    nx = x + dx
                    if nx < 0 or nx >= w:
                        continue
                    if alpha[ny][nx] == 255:
                        rr, gg, bb, _aa = px[nx, ny]
                        samples.append((rr, gg, bb))

            if samples:
                rr = sum(s[0] for s in samples) // len(samples)
                gg = sum(s[1] for s in samples) // len(samples)
                bb = sum(s[2] for s in samples) // len(samples)
                px[x, y] = (rr, gg, bb, 255)
            else:
                # fallback: keep own color, just solidify
                px[x, y] = (r, g, b, 255)

    return img


def apply_alpha_handling(img: Image.Image, mode: str, alpha_drop: int, bleed_radius: int) -> Image.Image:
    mode = mode.lower()
    if mode == "threshold_delete":
        return alpha_mode_threshold_delete(img, alpha_drop=alpha_drop)
    if mode == "solidify":
        return alpha_mode_solidify(img, alpha_drop=alpha_drop)
    if mode == "solidify_bleed":
        return alpha_mode_solidify_bleed(img, alpha_drop=alpha_drop, radius=bleed_radius)
    raise ValueError(f"Unknown --alpha-mode: {mode}")


def find_nontransparent_bounds(rgba: Image.Image) -> Optional[Tuple[int, int, int, int]]:
    if rgba.mode != "RGBA":
        rgba = rgba.convert("RGBA")

    px = rgba.load()
    w, h = rgba.size
    minx, miny = 10**9, 10**9
    maxx, maxy = -1, -1

    for y in range(h):
        for x in range(w):
            if px[x, y][3] != 0:
                if x < minx: minx = x
                if y < miny: miny = y
                if x > maxx: maxx = x
                if y > maxy: maxy = y

    if maxx < 0:
        return None
    return (minx, miny, maxx, maxy)


def encode_rle(frame_rgba: Image.Image, rle_data: bytearray) -> RLEHDR:
    if frame_rgba.mode != "RGBA":
        frame_rgba = frame_rgba.convert("RGBA")

    w, h = frame_rgba.size
    bounds = find_nontransparent_bounds(frame_rgba)
    dataptr = len(rle_data)

    if bounds is None:
        rle_data += b"\x00"
        return RLEHDR(0, 0, 0, 0, 0, 0, 2, dataptr)

    adjx, adjy, maxx, maxy = bounds
    wdh = (maxx - adjx + 1)
    hgh = (maxy - adjy + 1)

    hdr = RLEHDR(
        srcx=adjx, srcy=adjy,
        wdh=wdh, hgh=hgh,
        adjx=adjx, adjy=adjy,
        pixfmt=2,
        dataptr=dataptr
    )

    px = frame_rgba.load()

    transcount = 0
    oldx = 0
    currentrowlastx = 0
    did_color = False
    colarray: List[int] = []

    # Mirror POXStudio behaviour: iterate from adjy..height and adjx..width
    for y in range(adjy, h):
        for x in range(adjx, w):
            r, g, b, a = px[x, y]
            if a != 0:
                did_color = True

                if abs(transcount - oldx) > 0:
                    two = (transcount - oldx)
                    if two < 0:
                        two = two + adjx - 1
                    rl = int(two) << 1
                    rle_data += b"\x02" + struct.pack("<i", rl)  # signed
                    transcount = 0
                    oldx = 0

                colarray.append(color_to_rgb565(r, g, b))
                currentrowlastx = x
            else:
                if did_color:
                    did_color = False
                    rl = len(colarray)
                    rle_data += b"\x01" + struct.pack("<i", rl)  # signed
                    for c in colarray:
                        rle_data += struct.pack("<H", c)
                    colarray.clear()
                transcount += 1

        if did_color:
            did_color = False
            rl = len(colarray)
            rle_data += b"\x01" + struct.pack("<i", rl)  # signed
            for c in colarray:
                rle_data += struct.pack("<H", c)
            colarray.clear()

        rle_data += b"\x03"
        oldx = currentrowlastx
        transcount = 0

    rle_data += b"\x00"
    return hdr


def write_pox(out_path: str, res_type: str, ini_text: str, frames: List[Image.Image]) -> None:
    if len(res_type) != 2:
        raise ValueError("type must be exactly 2 chars (e.g. ST, CC, TT, ...)")
    res_type = res_type.upper()

    ini_bytes = ini_text.encode("cp1252", errors="replace")

    rle_data = bytearray()
    rle_hdrs: List[RLEHDR] = [encode_rle(fr, rle_data) for fr in frames]

    with open(out_path, "wb") as f:
        f.write(b"POXA")
        f.write(res_type.encode("ascii"))
        f.write(struct.pack("<H", 0x0A0D))
        f.write(struct.pack("<I", len(ini_bytes)))
        f.write(ini_bytes)

        f.write(struct.pack("<H", 0x4242))
        f.write(struct.pack("<I", len(frames)))
        f.write(struct.pack("<I", len(rle_data)))

        for hdr in rle_hdrs:
            f.write(hdr.pack())

        f.write(rle_data)
        f.write(struct.pack("<H", 0x4242))


def list_pngs(directory: str) -> List[str]:
    files = [os.path.join(directory, f) for f in os.listdir(directory) if f.lower().endswith(".png")]
    files.sort(key=natural_key)
    return files


def find_single_ini(directory: str) -> str:
    inis = [os.path.join(directory, f) for f in os.listdir(directory) if f.lower().endswith(".ini")]
    inis.sort(key=natural_key)
    if not inis:
        raise ValueError(f"No .ini found in: {directory}")
    if len(inis) > 1:
        raise ValueError(f"More than one .ini found in: {directory}")
    return inis[0]


def load_frames_from_dir(
    directory: str,
    key_rgb: Optional[Tuple[int, int, int]],
    alpha_mode: str,
    alpha_drop: int,
    bleed_radius: int
) -> List[Image.Image]:
    paths = list_pngs(directory)
    if not paths:
        raise ValueError(f"No PNGs found in: {directory}")

    frames: List[Image.Image] = []
    base_size = None

    for p in paths:
        im = Image.open(p)
        im = apply_keycolor_transparency(im, key_rgb)

        # IMPORTANT: alpha handling happens here (no changes to RLE logic)
        im = apply_alpha_handling(im, mode=alpha_mode, alpha_drop=alpha_drop, bleed_radius=bleed_radius)

        if base_size is None:
            base_size = im.size
        elif im.size != base_size:
            raise ValueError(
                f"Frame size mismatch in {directory}: expected {base_size}, got {im.size} at {p}"
            )
        frames.append(im)

    return frames


def generate_from_one_folder(folder: str, default_type: str, alpha_mode: str, alpha_drop: int, bleed_radius: int) -> None:
    folder = os.path.abspath(folder)
    ini_path = find_single_ini(folder)
    ini_text = read_text_cp1252(ini_path)

    res_type = extract_type_marker(ini_text) or default_type.upper()
    key_rgb = parse_transparent_color_from_ini(ini_text)

    frames = load_frames_from_dir(folder, key_rgb, alpha_mode=alpha_mode, alpha_drop=alpha_drop, bleed_radius=bleed_radius)
    if len(frames) > 999:
        raise ValueError(f"{folder}: frame count {len(frames)} exceeds 999 (requested limit).")

    ini_base = os.path.splitext(os.path.basename(ini_path))[0]
    parent_dir = os.path.dirname(folder)
    out_pox = os.path.join(parent_dir, ini_base + ".pox")

    tmp_pox = out_pox + ".tmp"
    write_pox(tmp_pox, res_type, ini_text, frames)
    os.replace(tmp_pox, out_pox)

    shutil.rmtree(folder, ignore_errors=False)


def generate_recursive(root: str, default_type: str, alpha_mode: str, alpha_drop: int, bleed_radius: int) -> None:
    root = os.path.abspath(root)

    for dirpath, _, filenames in os.walk(root):
        inis = [f for f in filenames if f.lower().endswith(".ini")]
        pngs = [f for f in filenames if f.lower().endswith(".png")]

        if len(inis) == 1 and len(pngs) >= 1:
            try:
                generate_from_one_folder(
                    dirpath,
                    default_type=default_type,
                    alpha_mode=alpha_mode,
                    alpha_drop=alpha_drop,
                    bleed_radius=bleed_radius
                )
                print(f"GENERATE OK: {dirpath}")
            except Exception as e:
                print(f"GENERATE FAIL: {dirpath} :: {e}")


# ----------------- CLI --------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(
        description="Single-file POX exporter/generator for Siege of Avalon."
    )

    group = ap.add_mutually_exclusive_group(required=True)
    group.add_argument("-export", dest="export_dir", help="Recursively export all POX files under DIR")
    group.add_argument("-generate", dest="generate_dir", help="Recursively build POX files from extracted folders under DIR")

    ap.add_argument("--default-type", default="ST", help="Fallback POX type if INI lacks ;POX_TOOL_TYPE=.. (default: ST)")

    ap.add_argument(
        "--alpha-mode",
        default="solidify_bleed",
        choices=["threshold_delete", "solidify", "solidify_bleed"],
        help="How to eliminate semi-transparent pixels during GENERATE. Default: solidify_bleed (recommended)."
    )
    ap.add_argument(
        "--alpha-drop",
        type=int,
        default=10,
        help="Alpha below this becomes fully transparent (0). Default: 10. For strict hard-threshold delete use 255."
    )
    ap.add_argument(
        "--bleed-radius",
        type=int,
        default=2,
        help="Neighbor radius used by solidify_bleed (RGB reconstruction). Default: 2."
    )

    args = ap.parse_args()

    if args.export_dir:
        export_recursive(args.export_dir)
    elif args.generate_dir:
        generate_recursive(
            args.generate_dir,
            default_type=args.default_type,
            alpha_mode=args.alpha_mode,
            alpha_drop=args.alpha_drop,
            bleed_radius=args.bleed_radius
        )


if __name__ == "__main__":
    main()