"""Build the Chinese UI font atlas + code tables for The I of the Dragon.

Approach
--------
The engine normalizes `Fonts.dat` char rects by the *actual* D3D texture size
(AddChar -> vtable+0x38 surface desc).  So if we replace the UI font texture
`UI\\Overhead\\Shrift_gb_Germany.TGA` (256x128) with a bigger atlas that still
contains the original ASCII artwork at pixel (0,0), every existing ASCII rect
keeps working unchanged.

CJK glyphs are rendered as full-width cells at the bottom/right of the atlas and
described by the side tables consumed by the `.cjk` code caves:

    cjk_glyph  : 65536 * float4 (u1,v1,u2,v2)
    cjk_metric : 65536 * int2   (advance, height)

`key = (lead_byte << 8) | trail_byte` -- exactly the code the cave builds.

Outputs (into --out):
    Shrift_gb_Germany.TGA   new texture (TGA type 2, 32bpp, bottom-up, same desc)
    cjk_glyph.bin
    cjk_metric.bin
    cells.json              code -> cell / unicode, for diagnostics
    atlas_preview.png       downscaled preview (not deployed)
"""

import argparse
import json
import os
import struct
import sys

from PIL import Image, ImageDraw, ImageFont

DEFAULT_GAME = r"F:\steam\steamapps\common\The I of the Dragon"
DEFAULT_OUT = r"E:\DragonHD\build\cjk"
TEX_NAME = r"UI\Overhead\Shrift_gb_Germany.TGA"
FONT_CANDIDATES = [
    (r"C:\Windows\Fonts\simsun.ttc", 0),
    (r"C:\Windows\Fonts\simhei.ttf", 0),
    (r"C:\Windows\Fonts\msyh.ttc", 0),
]


# ---------------------------------------------------------------- res reader
def read_res_entry(res_path, want_name):
    data = open(res_path, "rb").read()
    count = struct.unpack_from("<I", data, 0)[0]
    p = 4
    for _ in range(count):
        ln = struct.unpack_from("<I", data, p)[0]
        p += 4
        name = data[p:p + ln].decode("latin1")
        p += ln
        off, size = struct.unpack_from("<II", data, p)
        p += 8
        if name == want_name:
            return data[off:off + size]
    return None


def load_original_tga(game):
    loose = os.path.join(game, "Data", "Textures", "Ui", "Overhead",
                         os.path.basename(TEX_NAME.replace("\\", os.sep)))
    if os.path.exists(loose):
        return open(loose, "rb").read()
    blob = read_res_entry(os.path.join(game, "Data", "Textures.res"), TEX_NAME)
    if blob is None:
        raise SystemExit("cannot locate " + TEX_NAME)
    return blob


# ---------------------------------------------------------------- TGA <-> PIL
def tga_geometry(raw):
    if len(raw) < 18:
        raise SystemExit("not a TGA")
    idlen = raw[0]
    imgtype = raw[2]
    w = struct.unpack_from("<H", raw, 12)[0]
    h = struct.unpack_from("<H", raw, 14)[0]
    bpp = raw[16]
    desc = raw[17]
    return idlen, imgtype, w, h, bpp, desc


def tga_to_image(raw):
    idlen, imgtype, w, h, bpp, desc = tga_geometry(raw)
    if imgtype != 2 or bpp != 32:
        raise SystemExit("unsupported TGA type=%d bpp=%d" % (imgtype, bpp))
    off = 18 + idlen
    bottom_up = not (desc & 0x20)
    img = Image.new("RGBA", (w, h))
    dst = img.load()
    row_bytes = w * 4
    for y in range(h):
        src_row = (h - 1 - y) if bottom_up else y
        base = off + src_row * row_bytes
        row = raw[base:base + row_bytes]
        for x in range(w):
            b, g, r, a = row[x * 4:x * 4 + 4]
            dst[x, y] = (r, g, b, a)
    return img, desc


def image_to_tga(img, desc=0x08, footer=b""):
    w, h = img.size
    rgba = img.convert("RGBA")
    r, g, b, a = rgba.split()
    bgra = Image.merge("RGBA", (b, g, r, a))
    bottom_up = not (desc & 0x20)
    if bottom_up:
        bgra = bgra.transpose(Image.FLIP_TOP_BOTTOM)
    body = bgra.tobytes()
    hdr = (struct.pack("<BBB", 0, 0, 2)
           + struct.pack("<HHB", 0, 0, 0)
           + struct.pack("<HHHHBB", 0, 0, w, h, 32, desc))
    return hdr + body + footer


# ---------------------------------------------------------------- charset
def iter_gb2312():
    """code(int) -> unicode char, full GB2312 double byte space."""
    out = {}
    for hi in range(0xA1, 0xF8):
        for lo in range(0xA1, 0xFF):
            try:
                u = bytes([hi, lo]).decode("gb2312")
            except Exception:
                continue
            out[(hi << 8) | lo] = u
    return out


def chars_from_text(path, base):
    """merge chars found in a utf-8 text file into the base set."""
    txt = open(path, encoding="utf-8").read()
    used = {}
    for ch in txt:
        if ord(ch) < 0x80:
            continue
        try:
            enc = ch.encode("gb2312")
        except Exception:
            print("  ! cannot gb2312-encode %r" % ch, file=sys.stderr)
            continue
        if len(enc) != 2:
            continue
        used[(enc[0] << 8) | enc[1]] = ch
    return used


# ---------------------------------------------------------------- build
def pick_font(size):
    for path, idx in FONT_CANDIDATES:
        if os.path.exists(path):
            return ImageFont.truetype(path, size, index=idx), path
    raise SystemExit("no CJK ttf found")


def lower_full_stop(tile):
    """Anchor the Chinese full stop at the bottom of its full-width cell."""
    bbox = tile.getchannel('A').getbbox()
    if bbox is None:
        return tile.copy()
    result = Image.new('RGBA', tile.size, (0, 0, 0, 0))
    result.paste(tile.crop(bbox), (bbox[0], tile.height - (bbox[3] - bbox[1])))
    return result


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--game", default=DEFAULT_GAME)
    ap.add_argument("--out", default=DEFAULT_OUT)
    ap.add_argument("--size", type=int, default=12,
                    help="CJK pixel size (12 matches the UI pixel font)")
    ap.add_argument("--cell", type=int, default=12)
    ap.add_argument("--atlas-w", type=int, default=2048)
    ap.add_argument("--atlas-h", type=int, default=1024)
    ap.add_argument("--cjk-y0", type=int, default=128)
    ap.add_argument("--draw-dy", type=int, default=-1,
                    help="vertical offset of the glyph inside the cell")
    ap.add_argument("--charset", default="gb2312",
                    choices=["gb2312", "text"],
                    help="gb2312 = every double byte code; text = only chars in --text")
    ap.add_argument("--text", default=None)
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)

    orig_raw = load_original_tga(args.game)
    orig_img, desc = tga_to_image(orig_raw)
    ow, oh = orig_img.size
    footer = orig_raw[-26:] if len(orig_raw) >= 26 and orig_raw[-18:-1] == b"TRUEVISION-XFILE." else b""
    print("original %dx%d desc=0x%02x footer=%d" % (ow, oh, desc, len(footer)))

    base = iter_gb2312()
    if args.charset == "text":
        base = chars_from_text(args.text, base)
    print("charset size: %d" % len(base))

    cell = args.cell
    cols = args.atlas_w // cell
    rows = (args.atlas_h - args.cjk_y0) // cell
    if len(base) > cols * rows:
        raise SystemExit("charset %d does not fit in %dx%d (%d cells)"
                         % (len(base), cols, rows, cols * rows))

    atlas = Image.new("RGBA", (args.atlas_w, args.atlas_h), (0, 0, 0, 0))
    atlas.paste(orig_img, (0, 0))

    font, fontpath = pick_font(args.size)
    print("font: %s @ %dpx" % (fontpath, args.size))
    d = ImageDraw.Draw(atlas)

    cells = {}
    clipped = 0
    for i, code in enumerate(sorted(base)):
        uni = base[code]
        col = i % cols
        row = i // cols
        x = col * cell
        y = args.cjk_y0 + row * cell
        d.text((x, y + args.draw_dy), uni, font=font, fill=(255, 255, 255, 255))
        if uni == '\u3002':
            atlas.paste(lower_full_stop(atlas.crop((x, y, x + cell, y + cell))), (x, y))
        cells[code] = {"x": x, "y": y, "uni": uni}
        # sanity: ink must fit the cell
        bb = atlas.crop((x, y, x + cell, y + cell)).getbbox()
        if bb is None:
            clipped += 1

    print("placed %d glyphs (%d empty), grid %dx%d" % (len(cells), clipped, cols, rows))

    # tables
    glyph = bytearray(65536 * 16)
    metric = bytearray(65536 * 8)
    aw, ah = args.atlas_w, args.atlas_h
    for code, c in cells.items():
        x, y = c["x"], c["y"]
        struct.pack_into("<4f", glyph, code * 16,
                         x / aw, y / ah, (x + cell) / aw, (y + cell) / ah)
        struct.pack_into("<2i", metric, code * 8, cell, cell)

    open(os.path.join(args.out, "cjk_glyph.bin"), "wb").write(glyph)
    open(os.path.join(args.out, "cjk_metric.bin"), "wb").write(metric)

    tga = image_to_tga(atlas, desc=desc, footer=footer)
    open(os.path.join(args.out, "Shrift_gb_Germany.TGA"), "wb").write(tga)

    with open(os.path.join(args.out, "cells.json"), "w", encoding="utf-8") as f:
        json.dump({"cell": cell, "atlas": [aw, ah], "cjk_y0": args.cjk_y0,
                   "size": args.size, "font": fontpath, "cells": cells}, f)

    meta = {"cell": cell, "atlas": [aw, ah], "cjk_y0": args.cjk_y0,
            "size": args.size, "font": fontpath, "count": len(cells),
            "advance": cell, "height": cell}
    with open(os.path.join(args.out, "atlas_meta.json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)

    # preview
    prev = atlas.resize((args.atlas_w // 4, args.atlas_h // 4), Image.NEAREST)
    prev.save(os.path.join(args.out, "atlas_preview.png"))

    print("atlas TGA %d bytes, glyph/probe tables written to %s" % (len(tga), args.out))


if __name__ == "__main__":
    main()