import gzip
import json
import os
import sys

import numpy as np
from PIL import Image, PngImagePlugin

def byteize(alpha):
    alpha = alpha.T.reshape((-1,))
    alpha = alpha[: (alpha.shape[0] // 8) * 8]
    alpha = np.bitwise_and(alpha, 1)
    alpha = alpha.reshape((-1, 8))
    return np.packbits(alpha, axis=1)


class LSBExtractor:
    def __init__(self, alpha):
        self.data = byteize(alpha)
        self.pos = 0

    def get_next_n_bytes(self, n):
        n_bytes = self.data[self.pos : self.pos + n]
        self.pos += n
        return bytearray(n_bytes)

    def read_32bit_integer(self):
        bytes_list = self.get_next_n_bytes(4)
        if len(bytes_list) != 4:
            return None
        return int.from_bytes(bytes_list, byteorder="big")


def extract_stealth_png(path):
    with Image.open(path) as image:
        if "A" not in image.getbands():
            raise ValueError("image has no alpha channel")
        alpha = np.array(image.getchannel("A"))

    reader = LSBExtractor(alpha)
    magic = "stealth_pngcomp"
    read_magic = reader.get_next_n_bytes(len(magic)).decode("utf-8")
    if read_magic != magic:
        raise ValueError("magic number mismatch")

    read_len_bits = reader.read_32bit_integer()
    if read_len_bits is None:
        raise ValueError("missing payload length")
    read_len = read_len_bits // 8

    payload = reader.get_next_n_bytes(read_len)
    json_bytes = gzip.decompress(payload)
    try:
        json_data = json.loads(json_bytes.decode("utf-8"))
        if "Comment" in json_data and isinstance(json_data["Comment"], str):
            json_data["_CommentRaw"] = json_data["Comment"]
            json_data["Comment"] = json.loads(json_data["Comment"])
    except json.JSONDecodeError:
        return {"WebUI": json_bytes.decode("utf-8", errors="replace")}
    return json_data


def main() -> int:
    path = sys.argv[1] if len(sys.argv) > 1 else "21_138614363.png"
    if not os.path.exists(path):
        print(f"File not found: {path}")
        return 1

    data = extract_stealth_png(path)
    print(json.dumps(data, indent=2, ensure_ascii=True, sort_keys=True))
    if isinstance(data, dict) and "Comment" in data:
        comment_json = data.get("_CommentRaw")
        if not isinstance(comment_json, str):
            if isinstance(data["Comment"], str):
                comment_json = data["Comment"]
            else:
                comment_json = json.dumps(
                    data["Comment"],
                    ensure_ascii=True,
                    separators=(",", ":"),
                )
        reply = input("Create 1x1 metadata PNG? [y/N] ").strip().lower()
        if reply in ("y", "yes"):
            base, _ = os.path.splitext(path)
            out_path = f"{base}_meta.png"
            img = Image.new("RGBA", (1, 1), (0, 0, 0, 0))
            meta = PngImagePlugin.PngInfo()
            meta.add_text("Comment", comment_json)
            img.save(out_path, pnginfo=meta)
            print(f"Wrote {out_path}")
    input()
    return 0


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