import json
import struct
import sys
from pathlib import Path

# Base64 字母表（CC:Tweaked 使用的自定义顺序）
B64STR = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
B64MAP = {ch: i for i, ch in enumerate(B64STR)}


def custom_b64decode(s):
    """使用自定义字母表解码 base64 字符串（忽略 '='）"""
    val = 0
    bits = 0
    out = bytearray()
    for ch in s:
        if ch == '=':
            continue
        val = (val << 6) | B64MAP[ch]
        bits += 6
        if bits >= 8:
            bits -= 8
            out.append((val >> bits) & 0xFF)
    return bytes(out)


def decode_frame(data):
    # 验证头部
    assert data[:4] == b'\x00\x00\x00\x00', "Invalid frame header (first 4 bytes)"
    assert data[8:16] == b'\x00\x00\x00\x00\x00\x00\x00\x00', "Invalid frame header (bytes 8-15)"
    
    width, height = struct.unpack('<HH', data[4:8])
    total_chars = width * height

    # === 解码字符 RLE 流 ===
    text_bytes = []
    pos = 16
    remaining = total_chars
    while remaining > 0:
        char_byte = data[pos]
        count = data[pos + 1]
        pos += 2
        text_bytes.extend([char_byte] * count)
        remaining -= count
    assert len(text_bytes) == total_chars, "Character RLE length mismatch"

    # === 解码颜色 RLE 流 ===
    color_bytes = []
    remaining = total_chars
    while remaining > 0:
        color_byte = data[pos]
        count = data[pos + 1]
        pos += 2
        color_bytes.extend([color_byte] * count)
        remaining -= count
    assert len(color_bytes) == total_chars, "Color RLE length mismatch"

    # === 读取调色板（16 色 × RGB）===
    palette_raw = data[pos:pos + 48]
    assert len(palette_raw) == 48, "Palette must be 48 bytes"
    palette = []
    for i in range(16):
        r = palette_raw[i * 3]
        g = palette_raw[i * 3 + 1]
        b = palette_raw[i * 3 + 2]
        palette.append([r, g, b])

    # === 构建行数据 ===
    text_rows = []
    fg_rows = []
    bg_rows = []
    idx = 0
    for y in range(height):
        row_text_parts = []
        row_fg = ''
        row_bg = ''
        for x in range(width):
            b = text_bytes[idx]
            row_text_parts.append(f"\\{b:03d}")  # 如 \130

            c = color_bytes[idx]
            fg = c & 0x0F
            bg = (c >> 4) & 0x0F
            row_fg += f"{fg:x}"
            row_bg += f"{bg:x}"
            idx += 1
        text_rows.append(''.join(row_text_parts))
        fg_rows.append(row_fg)
        bg_rows.append(row_bg)

    return {
        "width": width,
        "height": height,
        "text": text_rows,
        "foreground": fg_rows,
        "background": bg_rows,
        "palette": palette
    }


def encode_frame(frame):
    width = frame["width"]
    height = frame["height"]
    total = width * height

    # === 从 \ddd 序列还原字节 ===
    full_text = ''.join(frame["text"])
    text_bytes = []
    i = 0
    while i < len(full_text):
        if full_text[i] == '\\' and i + 4 <= len(full_text):
            try:
                num_str = full_text[i+1:i+4]
                num = int(num_str)
                if 0 <= num <= 255:
                    text_bytes.append(num)
                    i += 4
                    continue
            except ValueError:
                pass
        raise ValueError(f"Invalid escape sequence at position {i}: expected \\ddd, got {full_text[i:i+5]}")
    assert len(text_bytes) == total, "Text byte count mismatch"

    # === 编码字符 RLE ===
    char_rle = []
    i = 0
    while i < total:
        b = text_bytes[i]
        j = i
        while j < total and text_bytes[j] == b and (j - i + 1) <= 255:
            j += 1
        count = j - i
        char_rle.extend([b, count])
        i = j

    # === 构建颜色流 ===
    color_stream = []
    for y in range(height):
        fg_row = frame["foreground"][y]
        bg_row = frame["background"][y]
        assert len(fg_row) == len(bg_row) == width, f"Row {y} length mismatch"
        for x in range(width):
            fg = int(fg_row[x], 16) & 0x0F
            bg = int(bg_row[x], 16) & 0x0F
            color_stream.append((bg << 4) | fg)
    assert len(color_stream) == total

    # === 编码颜色 RLE ===
    color_rle = []
    i = 0
    while i < total:
        c = color_stream[i]
        j = i
        while j < total and color_stream[j] == c and (j - i + 1) <= 255:
            j += 1
        count = j - i
        color_rle.extend([c, count])
        i = j

    # === 调色板 ===
    palette_bytes = []
    for rgb in frame["palette"]:
        assert len(rgb) == 3
        r, g, b = rgb
        assert 0 <= r <= 255 and 0 <= g <= 255 and 0 <= b <= 255
        palette_bytes.extend([r, g, b])
    assert len(palette_bytes) == 48

    # === 拼接二进制数据 ===
    header = b'\x00\x00\x00\x00' + struct.pack('<HH', width, height) + b'\x00' * 8
    body = bytes(char_rle) + bytes(color_rle) + bytes(palette_bytes)
    full_data = header + body

    # === 自定义 Base64 编码 ===
    b64 = ""
    val = 0
    bits = 0
    for byte in full_data:
        val = (val << 8) | byte
        bits += 8
        while bits >= 6:
            bits -= 6
            b64 += B64STR[(val >> bits) & 0x3F]
    if bits:
        b64 += B64STR[(val << (6 - bits)) & 0x3F]
    # 补齐到 4 的倍数（虽然 CC 可能不检查，但保持兼容）
    while len(b64) % 4:
        b64 += '='

    # 使用 !CPD 格式（12 位十六进制长度）
    hex_len = f"{len(b64):012x}"
    return f"!CPD{hex_len}{b64}"


def vid_to_json(vid_path, json_path):
    with open(vid_path, 'r', encoding='latin1') as f:
        lines = [line.rstrip('\n\r') for line in f]

    if not lines or lines[0] != "32Vid 1.1":
        raise ValueError("Not a valid .32vid file: missing header")

    fps_str = lines[1] if len(lines) > 1 else "0"
    try:
        fps = float(fps_str)
    except ValueError:
        fps = 0.0

    frame_lines = []
    for line in lines[2:]:
        if line.strip() == "":
            continue
        if line.startswith("!CP"):
            frame_lines.append(line)

    frames = []
    for line in frame_lines:
        mode = line[3]
        if mode == 'C':
            length = int(line[4:8], 16)
            b64data = line[8:8 + length]
        elif mode == 'D':
            length = int(line[4:16], 16)
            b64data = line[16:16 + length]
        else:
            raise ValueError(f"Unknown frame mode: {mode}")

        binary_data = custom_b64decode(b64data)
        frame = decode_frame(binary_data)
        frames.append(frame)

    result = {
        "header": {
            "magic": "32Vid 1.1",
            "fps": fps
        },
        "frames": frames
    }

    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(result, f, indent=2, ensure_ascii=False)


def json_to_vid(json_path, vid_path):
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    fps = data["header"]["fps"]
    frames = data["frames"]

    with open(vid_path, 'w', encoding='latin1') as f:
        f.write("32Vid 1.1\n")
        f.write(f"{fps}\n")
        for frame in frames:
            line = encode_frame(frame)
            f.write(line + "\n")


# ======================
# CLI 入口
# ======================
def main():
    if len(sys.argv) != 4:
        print("Usage:")
        print("  python 32vid_converter.py to-json input.32vid output.json")
        print("  python 32vid_converter.py to-vid input.json output.32vid")
        sys.exit(1)

    command = sys.argv[1]
    input_path = sys.argv[2]
    output_path = sys.argv[3]

    if command == "to-json":
        vid_to_json(input_path, output_path)
        print(f"✅ Converted: {input_path} → {output_path}")
    elif command == "to-vid":
        json_to_vid(input_path, output_path)
        print(f"✅ Converted: {input_path} → {output_path}")
    else:
        print(f"❌ Unknown command: {command}")
        sys.exit(1)


if __name__ == "__main__":
    main()