#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
i1_to_png.py — 将 LVGL 模拟器 `--screenshot` 产出的 I1(1bpp) 原始位图转换为 PNG。

输入格式（模拟器 main.c 的 --screenshot 模式输出）：
  - 1 位/像素，每字节 MSB 在左（bit7=该行最左像素），bit=1 为前景、bit=0 为背景
  - 每行 stride = 宽 / 8 字节，行内无填充
  - 模拟器当前已剥离 8 字节调色板，文件即纯像素数据（宽/8 × 高 字节）
  - 若遇到仍含调色板的旧文件（像素数据 + 8 字节），--auto-palette 会自动识别并跳过

用法示例：
  python i1_to_png.py shot.bin                       # -> shot.png（默认 160x160）
  python i1_to_png.py shot.bin -o out.png            # 指定输出
  python i1_to_png.py shot.bin --scale 6            # 放大 6 倍便于查看（最近邻，保持像素感）
  python i1_to_png.py shot.bin -W 320 -H 240        # 非默认分辨率
  python i1_to_png.py shot.bin --fg 0,0,0 --bg 255,255,255   # 自定义前/背景色（RGB 模式）
  python i1_to_png.py shot.bin --invert             # 前景/背景互换

依赖：Pillow（pip install pillow）。本仓库环境已验证 PIL 12.2.0。
退出码：0 成功；1 参数/数据错误。
"""

import argparse
import os
import sys

# DTU 屏幕默认分辨率（160x160，1bpp）
DEFAULT_W = 160
DEFAULT_H = 160
# I1 调色板大小（2 色 × 4 字节）。模拟器新版已剥离，旧文件可能残留。
I1_PALETTE_BYTES = 2 * 4


def _decode_i1(data, width, height):
    """把原始 I1 数据解成 (height, width) 的 0/1 矩阵。MSB 在左。"""
    stride = width // 8
    if width % 8 != 0:
        raise ValueError("宽度必须是 8 的整数倍（I1 按字节打包），当前 %d" % width)
    need = stride * height
    if len(data) < need:
        raise ValueError("数据不足：%d 字节 < 需要 %d 字节（%dx%d I1）"
                         % (len(data), need, width, height))
    px = []
    for y in range(height):
        base = y * stride
        row = [0] * width
        for x in range(width):
            byte_i = x >> 3
            bit_i = 7 - (x & 7)
            row[x] = (data[base + byte_i] >> bit_i) & 1
        px.append(row)
    return px


def _pick_palette_offset(data, width, height, auto):
    """返回像素数据在文件中的起始偏移（跳过调色板）。"""
    stride = width // 8
    need = stride * height
    sizes = [len(data)]
    if auto:
        # 像素数据前可能带 I1 调色板（8 字节）
        if len(data) == need + I1_PALETTE_BYTES:
            return I1_PALETTE_BYTES
    if len(data) < need:
        raise ValueError("数据不足：%d 字节 < 像素所需 %d 字节" % (len(data), need))
    return 0


def i1_to_png(input_path, output_path=None, width=DEFAULT_W, height=DEFAULT_H,
              scale=1, invert=False, auto_palette=True,
              fg=None, bg=None):
    """核心转换函数，可被其他脚本直接 import 调用。

    fg / bg 形如 (r,g,b)；为 None 时用黑白（L 模式）。
    返回输出文件路径。
    """
    from PIL import Image

    if not os.path.isfile(input_path):
        raise FileNotFoundError(input_path)

    with open(input_path, "rb") as f:
        data = f.read()

    off = _pick_palette_offset(data, width, height, auto_palette)
    px = _decode_i1(data[off:], width, height)

    nw, nh = width, height

    if fg is None and bg is None:
        # 灰度：bit=1 -> 白(255)，bit=0 -> 黑(0)
        img = Image.new("L", (nw, nh))
        ipx = img.load()
        for y in range(nh):
            for x in range(nw):
                bit = px[y][x]
                if invert:
                    bit ^= 1
                ipx[x, y] = 255 if bit else 0
    else:
        fgc = fg if fg is not None else (255, 255, 255)
        bgc = bg if bg is not None else (0, 0, 0)
        if invert:
            fgc, bgc = bgc, fgc
        img = Image.new("RGB", (nw, nh), bgc)
        ipx = img.load()
        for y in range(nh):
            for x in range(nw):
                bit = px[y][x]
                ipx[x, y] = fgc if bit else bgc

    # 放大（最近邻，保持 1bpp 像素感）
    if scale and scale > 1:
        img = img.resize((nw * scale, nh * scale), Image.NEAREST)

    if output_path is None:
        output_path = os.path.splitext(input_path)[0] + ".png"
    img.save(output_path, "PNG")
    return output_path


def _parse_rgb(s):
    if s is None:
        return None
    parts = s.split(",")
    if len(parts) != 3:
        raise argparse.ArgumentTypeError("颜色格式应为 R,G,B，如 255,0,0")
    try:
        r, g, b = (int(p) for p in parts)
    except ValueError:
        raise argparse.ArgumentTypeError("颜色分量必须是整数")
    for v in (r, g, b):
        if not (0 <= v <= 255):
            raise argparse.ArgumentTypeError("颜色分量需在 0~255")
    return (r, g, b)


def _ensure_utf8_console():
    """Windows 控制台默认 GBK，中文会乱码；统一切到 UTF-8（不支持则静默跳过）。"""
    for stream in (sys.stdout, sys.stderr):
        try:
            stream.reconfigure(encoding="utf-8")
        except (AttributeError, ValueError):
            pass


def main(argv=None):
    _ensure_utf8_console()
    ap = argparse.ArgumentParser(
        description="将模拟器 I1(1bpp) 原始截图转换为 PNG",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    ap.add_argument("input", help="I1 原始位图文件（.bin）")
    ap.add_argument("-o", "--output", default=None,
                    help="输出 PNG 路径（默认与输入同名 .png）")
    ap.add_argument("-W", "--width", type=int, default=DEFAULT_W, help="图像宽度")
    ap.add_argument("-H", "--height", type=int, default=DEFAULT_H, help="图像高度")
    ap.add_argument("--scale", type=int, default=1,
                    help="放大倍数（最近邻，便于查看小图）")
    ap.add_argument("--invert", action="store_true", help="前景/背景互换")
    ap.add_argument("--no-auto-palette", action="store_true",
                    help="不做调色板自动识别（按纯像素数据解析）")
    ap.add_argument("--fg", type=_parse_rgb, default=None, help="前景色 R,G,B（bit=1）")
    ap.add_argument("--bg", type=_parse_rgb, default=None, help="背景色 R,G,B（bit=0）")
    args = ap.parse_args(argv)

    try:
        out = i1_to_png(
            args.input,
            output_path=args.output,
            width=args.width,
            height=args.height,
            scale=args.scale,
            invert=args.invert,
            auto_palette=not args.no_auto_palette,
            fg=args.fg,
            bg=args.bg,
        )
    except (ValueError, FileNotFoundError) as e:
        print("错误：%s" % e, file=sys.stderr)
        return 1

    print("已生成：%s" % out)
    return 0


if __name__ == "__main__":
    sys.exit(main())