#!/usr/bin/env python3
"""fb2png.py —— 把 tb_saolei 导出的显存 dump 还原成 PNG（等价于 VGA 截图）。

用法:
    py saolei/tools/fb2png.py <dump.txt> [输出.png]

dump 格式（tb_saolei.v 里的 dump_fb 任务生成）:
    CURSOR <row> <col>
    CTRL 0x<byte>            # [0] 光标可见  [6:4] 背景色
    FB 1200                  # 随后 1200 行十六进制字节，按格号 0..1199

渲染依据（与 RTL 严格对应，改 RTL 时这两个表要同步）:
  * 字库 cpu21_vga.v 的 font_row()：8x8、bit7 是最左像素，放大 2 倍 = 16x16 一格
  * 调色板 cpu21_vga.v 的 pal()：3 位 -> 4:4:4 RGB（0xA 分量 = 170）
  * 光标 cpu21_vga.v 的 cursor_box：格子四边一圈（1 像素宽）显示浅灰
  * 画面 40x30 格 -> 640x480

另外还会在终端打印一份 ASCII 图，方便在没有图片查看器的地方快速看结果。
"""

import struct
import sys
import zlib

COLS, ROWS = 40, 30
CELL_PX = 16            # 8x8 字模放大 2 倍

# --- 字库（8 行/字，bit7 = 最左），抄自 cpu21_vga.v 的 font_row() ---
FONT = {
    0:  [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],  # 空白
    1:  [0x18, 0x38, 0x18, 0x18, 0x18, 0x18, 0x7E, 0x00],  # 1
    2:  [0x3C, 0x66, 0x06, 0x0C, 0x18, 0x30, 0x7E, 0x00],  # 2
    3:  [0x3C, 0x66, 0x06, 0x1C, 0x06, 0x66, 0x3C, 0x00],  # 3
    4:  [0x0C, 0x1C, 0x3C, 0x6C, 0x7E, 0x0C, 0x0C, 0x00],  # 4
    5:  [0x7E, 0x60, 0x7C, 0x06, 0x06, 0x66, 0x3C, 0x00],  # 5
    6:  [0x1C, 0x30, 0x60, 0x7C, 0x66, 0x66, 0x3C, 0x00],  # 6
    7:  [0x7E, 0x66, 0x0C, 0x18, 0x18, 0x18, 0x18, 0x00],  # 7
    8:  [0x3C, 0x66, 0x66, 0x3C, 0x66, 0x66, 0x3C, 0x00],  # 8
    9:  [0xFF, 0x81, 0xBD, 0xA5, 0xA5, 0xBD, 0x81, 0xFF],  # 未揭开
    10: [0x40, 0x70, 0x78, 0x70, 0x40, 0x40, 0x40, 0x00],  # 旗
    11: [0x3C, 0x7E, 0xFF, 0xFF, 0xFF, 0xFF, 0x7E, 0x3C],  # 雷
    12: [0x81, 0x42, 0x24, 0x18, 0x18, 0x24, 0x42, 0x81],  # 爆炸
    13: [0x3C, 0x66, 0x06, 0x0C, 0x18, 0x00, 0x18, 0x00],  # 问号
    14: [0x00, 0x00, 0x00, 0xFF, 0xFF, 0x00, 0x00, 0x00],  # 横线
    15: [0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF],  # 实心
}

# --- 调色板（3 位 -> RGB），抄自 cpu21_vga.v 的 pal()，4 位分量按 *17 扩到 8 位 ---
PAL = [
    (0x0, 0x0, 0x0),   # 0 黑
    (0x0, 0x0, 0xA),   # 1 深蓝
    (0x0, 0xA, 0x0),   # 2 绿
    (0x0, 0xA, 0xA),   # 3 青
    (0xA, 0x0, 0x0),   # 4 红
    (0xA, 0x0, 0xA),   # 5 洋红
    (0xA, 0xA, 0x0),   # 6 黄
    (0xA, 0xA, 0xA),   # 7 浅灰
]

GLYPH_CH = {0: ' ', 9: '#', 10: 'P', 11: '*', 12: 'X', 13: '?', 14: '-', 15: '='}
COLOR_CH = {0: '.', 1: 'b', 2: 'g', 3: 'c', 4: 'r', 5: 'm', 6: 'y', 7: 'w'}


def write_png(path, width, height, rows):
    """rows: 每行一个 bytes（长度 width*3），最小 PNG 编码。"""
    raw = b"".join(b"\x00" + r for r in rows)

    def chunk(tag, data):
        return (struct.pack(">I", len(data)) + tag + data
                + struct.pack(">I", zlib.crc32(tag + data) & 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(raw, 9))
    png += chunk(b"IEND", b"")
    with open(path, "wb") as f:
        f.write(png)


def parse_dump(path):
    cur = (None, None)
    ctrl = 0
    fb = []
    with open(path, encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if line.startswith("CURSOR"):
                p = line.split()
                cur = (int(p[1]), int(p[2]))
            elif line.startswith("CTRL"):
                ctrl = int(line.split()[1], 16)
            elif line.startswith("FB"):
                fb = []
            elif len(line) == 2 and fb is not None:
                try:
                    fb.append(int(line, 16))
                except ValueError:
                    pass
    if len(fb) != 1200:
        raise SystemExit("显存字节数不对：%d（应为 1200）" % len(fb))
    return cur, ctrl, fb


def main():
    if len(sys.argv) < 2:
        raise SystemExit(__doc__)
    src = sys.argv[1]
    dst = sys.argv[2] if len(sys.argv) > 2 else src.rsplit(".", 1)[0] + ".png"
    cur, ctrl, fb = parse_dump(src)

    bg = (ctrl >> 4) & 0x7
    cursor_box = bool(ctrl & 1) and cur[0] is not None
    print("光标 = (%s, %s)   光标可见 = %s   背景色 = %d"
          % (cur[0], cur[1], bool(ctrl & 1), bg))

    # ---- 渲染 ----
    rows = []
    for y in range(ROWS * CELL_PX):
        row = bytearray()
        cy, fy = y // CELL_PX, (y % CELL_PX) // 2
        for x in range(COLS * CELL_PX):
            cx, fx = x // CELL_PX, (x % CELL_PX) // 2
            cell = fb[cy * COLS + cx]
            g, c = cell & 0xF, (cell >> 4) & 0x7
            px_on = (FONT[g][fy] >> (7 - fx)) & 1
            col = c if px_on else bg
            if cursor_box and (cx, cy) == cur:
                # 格子四边一圈（1 像素）显示浅灰
                ox, oy = x % CELL_PX, y % CELL_PX
                if ox == 0 or ox == CELL_PX - 1 or oy == 0 or oy == CELL_PX - 1:
                    col = 7
            r, gg, b = PAL[col]
            row += bytes((r * 17, gg * 17, b * 17))
        rows.append(bytes(row))
    write_png(dst, COLS * CELL_PX, ROWS * CELL_PX, rows)
    print("已生成 %s（%dx%d）" % (dst, COLS * CELL_PX, ROWS * CELL_PX))

    # ---- 终端 ASCII 图 ----
    print("+" + "-" * COLS + "+")
    for y in range(ROWS):
        line = ""
        for x in range(COLS):
            cell = fb[y * COLS + x]
            g, c = cell & 0xF, (cell >> 4) & 0x7
            if cursor_box and (x, y) == cur:
                line += "["
            elif g in GLYPH_CH:
                line += GLYPH_CH[g]
            else:
                line += chr(ord('0') + g)      # 字形 1..8 就是数字 1..8
        line = line.replace("[", "[]" * 1)     # 光标只提示左边界，避免打乱列数
        print("|" + (line + " " * COLS)[:COLS] + "|")
    print("+" + "-" * COLS + "+")
    print("图例： #=未揭开  P=旗  *=雷  X=爆炸  ?=问号  -=横线  ==实心  1..8=数字")


main()