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

用法:
    py saolei/tools/fb2png.py <dump.txt> [输出.png]
    py saolei/tools/fb2png.py --term <dump.txt>      # 只在终端打印彩色棋盘（交互玩扫雷用）
    py saolei/tools/fb2png.py --quiet <dump.txt> [输出.png]   # 只写 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：格子四边一圈 2 像素宽的黄色框（常亮）
  * 画面 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 cell_char(g):
    """字形号 -> 一个可见字符（字形 1..8 就是数字 1..8）。"""
    if g in GLYPH_CH:
        return GLYPH_CH[g]
    if 1 <= g <= 8:
        return chr(ord('0') + g)
    return ' '


def dump_note(path):
    """把 dump 开头的 TIME / CURSOR / STATS 拼成一行，当终端画面的标题。"""
    note = []
    try:
        with open(path, encoding="utf-8", errors="replace") as f:
            for line in f:
                if line.startswith(("TIME", "CURSOR", "STATS")):
                    note.append(line.strip())
                elif line.startswith("FB"):
                    break
    except OSError:
        pass
    return "   ".join(note)


def draw_term(fb, cur, cursor_box, note=""):
    """在终端用 ANSI 24 位颜色画 40x30 棋盘（每格一个字符），光标格反色。"""
    try:                                    # Windows 控制台默认不是 UTF-8
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass
    out = ["\x1b[H\x1b[2J"]                 # 回左上角 + 清屏（连续刷新不闪屏残留）
    if note:
        out.append("  " + note + "\n")
    out.append("    " + "".join(str(x % 10) for x in range(COLS)) + "\n")
    for y in range(ROWS):
        cells = []
        for x in range(COLS):
            cell = fb[y * COLS + x]
            g, c = cell & 0xF, (cell >> 4) & 0x7
            r, gg, b = PAL[c]
            ch = cell_char(g)
            if cursor_box and (x, y) == cur:
                cells.append("\x1b[7m\x1b[38;2;%d;%d;%dm%s" % (r * 17, gg * 17, b * 17,
                                                              ch if ch != ' ' else '_'))
            else:
                cells.append("\x1b[38;2;%d;%d;%dm%s" % (r * 17, gg * 17, b * 17, ch))
        out.append("%2d  " % y + "".join(cells) + "\x1b[0m\n")
    out.append("\x1b[0m")
    out.append("图例： #=未揭开  P=旗  *=雷  X=爆炸  ?=问号  1..8=数字  空格=已揭开的空地；反色格 = 光标\n")
    sys.stdout.write("".join(out))
    sys.stdout.flush()


def main():
    args = list(sys.argv[1:])
    term_only = "--term" in args
    quiet = "--quiet" in args
    if term_only:
        args.remove("--term")
    if quiet:
        args.remove("--quiet")
    if quiet:                               # 只出图不说话（给脚本批量截图用）
        import os
        sys.stdout = open(os.devnull, "w")
    if not args:
        raise SystemExit(__doc__)
    src = args[0]
    dst = args[1] if len(args) > 1 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

    if term_only:
        draw_term(fb, cur, cursor_box, dump_note(src))
        return

    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:
                # 光标格：四边 2 像素黄色框，与 cpu21_vga.v 的 cursor_box 一致
                ox, oy = x % CELL_PX, y % CELL_PX
                if ox < 2 or ox > CELL_PX - 3 or oy < 2 or oy > CELL_PX - 3:
                    col = 6
            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 图（彩色版；PNG 已经写好了，这里只是方便快速看一眼）----
    draw_term(fb, cur, cursor_box, dump_note(src))


main()