#!/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 2400                  # 随后 2400 行十六进制字节：每格 2 字节（字形、颜色）

渲染依据（与 RTL 严格对应）:
  * 字库：直接从 cpu21_vga.v 的 font_row() 解析（8x8、bit7 最左，放大 2 倍 = 16x16 一格）
    —— RTL 改了字库（比如新增 A-Z）这里会自动跟随，不用手工抄表
  * 调色板 cpu21_vga.v 的 pal()：3 位 -> 4:4:4 RGB（0xA 分量 = 170）
  * 显存每格 2 字节：byte0 = 字形[5:0]，byte1 = 颜色[2:0]
  * 光标 cpu21_vga.v 的 cursor_box：格子四边一圈 2 像素宽的黄色框（常亮）
  * 画面 40x30 格 -> 640x480

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

import os
import re
import struct
import sys
import zlib

COLS, ROWS = 40, 30
CELL_PX = 16            # 8x8 字模放大 2 倍
FB_BYTES = COLS * ROWS * 2      # 每格 2 字节（字形 + 颜色）= 2400

RTL_VGA = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                        "..", "..", "pipelined-cpu", "srcs", "rtl",
                                        "cpu21_vga.v"))


def load_font_from_rtl(path=None):
    """直接从 cpu21_vga.v 的 font_row() 里解析 8x8 字模，返回 {字形号: [8 行]}。

    RTL 里每个字形形如
        6'd16: font_row = (r == 3'd0) ? 8'h7C : (r == 3'd1) ? 8'h66 : ... ;
    或全部相同行时形如
        6'd15: font_row = 8'hFF;
    两种情况都能解析。
    """
    path = path or RTL_VGA
    try:
        with open(path, encoding="utf-8", errors="replace") as f:
            text = f.read()
    except OSError as e:
        raise SystemExit("读不到 RTL 字库 %s：%s" % (path, e))
    m = re.search(r"function\s*\[7:0\]\s*font_row.*?endfunction", text, re.S)
    if not m:
        raise SystemExit("%s 里找不到 font_row()" % path)

    font = {}
    parts = re.split(r"6'd(\d+)\s*:", m.group(0))
    for i in range(1, len(parts) - 1, 2):       # 奇数下标是字形号，后一格是它的分支体
        idx, seg = int(parts[i]), parts[i + 1]
        rows = [0] * 8
        pairs = re.findall(r"\(r == 3'd(\d+)\)\s*\?\s*8'h([0-9A-Fa-f]{2})", seg)
        if pairs:
            for k, v in pairs:
                rows[int(k)] = int(v, 16)
        else:                                    # 形如 "font_row = 8'hFF;"：8 行一样
            one = re.search(r"font_row\s*=\s*8'h([0-9A-Fa-f]{2})", seg)
            if one:
                rows = [int(one.group(1), 16)] * 8
        font[idx] = rows
    if not font:
        raise SystemExit("%s 的 font_row() 没解析出任何字形" % path)
    return font

# --- 调色板（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) != FB_BYTES:
        raise SystemExit("显存字节数不对：%d（应为 %d；新版每格 2 字节）"
                         % (len(fb), FB_BYTES))
    return cur, ctrl, fb


def cell_of(fb, idx):
    """第 idx 格 -> (字形, 颜色)：byte0 = 字形[5:0]，byte1 = 颜色[2:0]。"""
    return fb[2 * idx] & 0x3F, fb[2 * idx + 1] & 0x07


def cell_char(g):
    """字形号 -> 一个可见字符（1..8 数字，16..41 = A-Z，其余见 GLYPH_CH）。"""
    if 16 <= g <= 41:
        return chr(ord('A') + g - 16)
    if 1 <= g <= 8:
        return str(g)
    return GLYPH_CH.get(g, '?')


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):
            g, c = cell_of(fb, y * COLS + x)
            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=爆炸  A-Z=文字  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:                               # 只出图不说话（给脚本批量截图用）
        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)
    font = load_font_from_rtl()

    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
            g, c = cell_of(fb, cy * COLS + cx)
            px_on = (font.get(g, [0] * 8)[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()