#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Phase C — vendor palette extraction from libcxsdk.so.

Scans the official native library for static 256-entry ARGB palette tables and
compares any candidates against the iron-bow table the project already owns
(OfficialTables.PALETTE256_ARGB), which is the extraction anchor.

RESULT ON THE SHIPPED LIBRARY (recorded in
analysis/sdk_re/android_app/palette_extraction_findings.md): the scan finds NO
static tables — libcxsdk.so builds all 12 palettes at runtime with arithmetic in
CFunctions::SetColorPalette. This script is therefore kept as the reproducible
proof of that negative (and would catch a table if a future vendor build embeds
one). The actual recovery was done by porting that routine; see
analysis/tools/PalIdentify.java + analysis/tools/PalExport2.java.

Pure standard library (struct + zlib + json). Writes UTF-8 explicitly.

Usage:
    python3 extract_palettes.py <libcxsdk.so> <OfficialTables.kt> <out_dir>
"""

import json
import os
import re
import struct
import sys
import zlib

ALPHA_RUN = 240          # minimum entries with alpha==0xFF and non-zero RGB
ENTRIES = 256
RUN_BYTES = ENTRIES * 4


def read_anchor(kt_path):
    """Extract PALETTE256_ARGB (256 ARGB ints) from OfficialTables.kt."""
    with open(kt_path, "r", encoding="utf-8") as fh:
        src = fh.read()
    m = re.search(r"val PALETTE256_ARGB = intArrayOf\((.*?)\)", src, re.S)
    if not m:
        raise SystemExit("PALETTE256_ARGB not found in " + kt_path)
    vals = [int(x) for x in re.findall(r"-?\d+", m.group(1))]
    if len(vals) != ENTRIES:
        raise SystemExit("anchor has %d entries, expected %d" % (len(vals), ENTRIES))
    return vals


def u32le(buf, off):
    return struct.unpack_from("<I", buf, off)[0]


def scan(data):
    """Return list of offsets where a 256-entry alpha-0xFF run starts."""
    hits = []
    limit = len(data) - RUN_BYTES
    for off in range(0, limit + 1, 4):
        ok = 0
        for i in range(0, ENTRIES, 16):          # sample every 16th entry first
            v = u32le(data, off + i * 4)
            if (v >> 24) == 0xFF and (v & 0xFFFFFF) != 0:
                ok += 1
        if ok == ENTRIES // 16:                   # cheap pass, verify fully
            full = 0
            for i in range(ENTRIES):
                v = u32le(data, off + i * 4)
                if (v >> 24) == 0xFF and (v & 0xFFFFFF) != 0:
                    full += 1
            if full >= ALPHA_RUN:
                hits.append(off)
    return hits


def merge(hits):
    """Merge overlapping runs (a 1024 B table can start at off, off+4, ...)."""
    merged = []
    for off in hits:
        if merged and off <= merged[-1][1]:
            merged[-1][1] = max(merged[-1][1], off + RUN_BYTES)
        else:
            merged.append([off, off + RUN_BYTES])
    return merged


def write_png(path, rows, row_h, scale):
    """Minimal PNG writer: one horizontal strip per row of 256 colours."""
    width = ENTRIES * scale
    height = len(rows) * row_h
    raw = bytearray()
    for row in rows:
        line = bytearray()
        for x in range(width):
            argb = row[x // scale]
            line += bytes(((argb >> 16) & 0xFF, (argb >> 8) & 0xFF, argb & 0xFF))
        for _ in range(row_h):
            raw.append(0)             # filter type 0
            raw += line

    def chunk(tag, payload):
        return (struct.pack(">I", len(payload)) + tag + payload
                + struct.pack(">I", zlib.crc32(tag + payload) & 0xFFFFFFFF))

    png = b"\x89PNG\r\n\x1a\n"
    png += chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
    png += chunk(b"IDAT", zlib.compress(bytes(raw), 9))
    png += chunk(b"IEND", b"")
    with open(path, "wb") as fh:
        fh.write(png)


def main():
    if len(sys.argv) != 4:
        raise SystemExit(__doc__)
    so_path, kt_path, out_dir = sys.argv[1:4]
    os.makedirs(out_dir, exist_ok=True)

    with open(so_path, "rb") as fh:
        data = fh.read()
    anchor = read_anchor(kt_path)
    print("scanned %s: %d bytes; anchor %d entries" % (so_path, len(data), len(anchor)))

    hits = scan(data)
    print("static 256-entry alpha=0xFF runs: %d" % len(hits))
    runs = merge(hits)

    candidates = []
    anchor_at = None
    for start, end in runs:
        for off in range(start, end - RUN_BYTES + 1, 4):
            if all(u32le(data, off + i * 4) == anchor[i] for i in range(ENTRIES)):
                anchor_at = off
                break
        candidates.append({
            "offset": start,
            "offsetHex": "0x%X" % start,
            "hexPreview": ["0x%08X" % u32le(data, start + i * 4) for i in range(16)],
            "length": ENTRIES,
        })

    if anchor_at is not None:
        print("ironbow anchor found at 0x%X" % anchor_at)
    else:
        print("ironbow anchor NOT found — palettes are not stored as static tables")

    report = {
        "binary": os.path.basename(so_path),
        "binarySize": len(data),
        "anchorFound": anchor_at is not None,
        "anchorOffset": anchor_at,
        "staticCandidates": candidates,
        "conclusion": ("no static palette tables: CFunctions::SetColorPalette builds "
                       "them at runtime" if anchor_at is None else "static tables present"),
    }
    json_path = os.path.join(out_dir, "palette_candidates.json")
    with open(json_path, "w", encoding="utf-8") as fh:
        json.dump(report, fh, ensure_ascii=False, indent=2)
    print("wrote %s" % json_path)

    # strip sheet: one row per candidate (empty sheet when none are found)
    rows = []
    for start, end in runs:
        rows.append([u32le(data, start + i * 4) for i in range(ENTRIES)])
    if not rows:
        rows.append([0xFF000000] * ENTRIES)     # placeholder so the file is valid
    png_path = os.path.join(out_dir, "palette_candidates.png")
    write_png(png_path, rows, row_h=32, scale=8)
    print("wrote %s (%d rows)" % (png_path, len(rows)))


if __name__ == "__main__":
    main()