"""Deploy / restore the Chinese font patch into the game directory.

`install`:
  1. back up the pristine exe / Textures.res / Strings.dat (once, into backup_game)
  2. copy the generated `.cjk` tables, rebuild the patched exe (apply_cjk)
  3. drop the patched exe into the game folder
  4. install the font atlas (loose file + inside Textures.res)
  5. patch Strings.dat with the Chinese `English` column

`restore`:
  puts every original file back and removes anything we added.
"""

import argparse
import hashlib
import json
import os
import shutil
import struct
import sys

GAME = r"F:\steam\steamapps\common\The I of the Dragon"
BUILD = r"E:\DragonHD\build\cjk"
PATCH = r"E:\DragonHD\patch"
BACKUP = r"E:\DragonHD\backup_game"

EXE = "TheIOfTheDragon.exe"
ATLAS_RES_NAME = r"UI\Overhead\Shrift_gb_Germany.TGA"
ATLAS_LOOSE = os.path.join("Data", "Textures", "Ui", "Overhead", "Shrift_gb_Germany.TGA")

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import apply_cjk  # noqa: E402
import make_cn_strings  # noqa: E402
from cn_strings import CN_STRINGS  # noqa: E402


def sha1(path):
    h = hashlib.sha1()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def backup_once(src, dst):
    if os.path.exists(dst):
        return False
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    shutil.copy2(src, dst)
    return True


def replace_res_entries(res_path, repl):
    data = open(res_path, "rb").read()
    count = struct.unpack_from("<I", data, 0)[0]
    p = 4
    entries = []
    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
        entries.append([name, data[off:off + size]])

    hdr = 4
    for name, _ in entries:
        hdr += 4 + len(name) + 8
    data_start = (hdr + 15) & ~15

    out = bytearray(struct.pack("<I", len(entries)))
    index = bytearray()
    payload = bytearray()
    offset = data_start
    for i, (name, blob) in enumerate(entries):
        blob = repl.get(name, blob)
        nb = name.encode("latin1")
        index += struct.pack("<I", len(nb)) + nb + struct.pack("<II", offset, len(blob))
        payload += blob
        if i != len(entries) - 1:
            pad = ((len(blob) + 15) & ~15) - len(blob)
            payload += b"\x00" * pad
            offset += len(blob) + pad
        else:
            offset += len(blob)
    out += index
    out += b"\x00" * (data_start - len(out))
    out += payload
    return bytes(out)


def install(args):
    os.makedirs(BACKUP, exist_ok=True)
    state_path = os.path.join(BACKUP, "state.json")
    state = {}
    if os.path.exists(state_path):
        state = json.load(open(state_path))

    game_exe = os.path.join(args.game, EXE)
    game_res = os.path.join(args.game, "Data", "Textures.res")
    game_str = os.path.join(args.game, "Data", "Misc", "Strings.dat")
    game_loose = os.path.join(args.game, ATLAS_LOOSE)

    def bak(name):
        return os.path.join(BACKUP, name)

    for src, dst in [(game_exe, bak(EXE)),
                     (game_res, bak("Textures.res")),
                     (game_str, bak("Strings.dat"))]:
        if backup_once(src, dst):
            print("backup", os.path.basename(dst))

    if "loose_existed" not in state:
        state["loose_existed"] = os.path.exists(game_loose)
        if state["loose_existed"]:
            backup_once(game_loose, bak("Shrift_gb_Germany.TGA"))
        json.dump(state, open(state_path, "w"))

    # 1. tables + patched exe (build against the pristine exe)
    os.makedirs(os.path.join(PATCH, "tables"), exist_ok=True)
    for t in ("cjk_glyph.bin", "cjk_metric.bin"):
        shutil.copy2(os.path.join(BUILD, t), os.path.join(PATCH, "tables", t))
    patch_out = os.path.join(PATCH, "TheIOfTheDragon_cjk.exe")
    apply_cjk.main(bak(EXE), patch_out, os.path.join(PATCH, "tables"))
    shutil.copy2(patch_out, game_exe)
    print("patched exe ->", game_exe)

    # 2. atlas: loose + res
    atlas_src = os.path.join(BUILD, "Shrift_gb_Germany.TGA")
    atlas_blob = open(atlas_src, "rb").read()
    os.makedirs(os.path.dirname(game_loose), exist_ok=True)
    open(game_loose, "wb").write(atlas_blob)
    print("loose atlas ->", game_loose)

    if not args.skip_res:
        open(game_res, "wb").write(
            replace_res_entries(bak("Textures.res"), {ATLAS_RES_NAME: atlas_blob}))
        print("res atlas updated ->", game_res)

    # 3. strings
    raw = open(bak("Strings.dat"), "rb").read()
    out, done, missing = make_cn_strings.patch(raw, CN_STRINGS)
    open(game_str, "wb").write(out)
    print("strings patched %d/%d" % (len(done), len(CN_STRINGS)))
    if missing:
        print("MISSING:", missing, file=sys.stderr)
        sys.exit(2)

    print("install done. exe sha1", sha1(game_exe), "res", sha1(game_res))


def restore(args):
    game_exe = os.path.join(args.game, EXE)
    game_res = os.path.join(args.game, "Data", "Textures.res")
    game_str = os.path.join(args.game, "Data", "Misc", "Strings.dat")
    game_loose = os.path.join(args.game, ATLAS_LOOSE)
    state_path = os.path.join(BACKUP, "state.json")
    state = json.load(open(state_path)) if os.path.exists(state_path) else {}

    for name, dst in [(EXE, game_exe), ("Textures.res", game_res), ("Strings.dat", game_str)]:
        src = os.path.join(BACKUP, name)
        if os.path.exists(src):
            shutil.copy2(src, dst)
            print("restored", dst)

    if os.path.exists(game_str + ".orig"):
        os.remove(game_str + ".orig")

    if state.get("loose_existed"):
        src = os.path.join(BACKUP, "Shrift_gb_Germany.TGA")
        if os.path.exists(src):
            shutil.copy2(src, game_loose)
            print("restored loose atlas")
    elif os.path.exists(game_loose):
        os.remove(game_loose)
        print("removed loose atlas")

    print("restore done. exe sha1", sha1(game_exe), "res", sha1(game_res))
    print("strings sha1", sha1(game_str))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("action", choices=["install", "restore"])
    ap.add_argument("--game", default=GAME)
    ap.add_argument("--skip-res", action="store_true")
    args = ap.parse_args()
    if args.action == "install":
        install(args)
    else:
        restore(args)


if __name__ == "__main__":
    main()