#!/usr/bin/env python3
"""Display or edit a Hengband SAVE26 character's height and weight."""

from __future__ import annotations

import argparse
import shutil
import struct
import sys
from pathlib import Path


VARIANT = b"Hengband"
QUICK_START_FIXED_PREFIX = 160
SUPPORTED_SAVE_VERSIONS = range(23, 27)


def cstring(data: bytearray, pos: int, limit: int = 4096) -> tuple[bytes, int]:
    end = data.find(0, pos, min(len(data), pos + limit))
    if end < 0:
        raise ValueError("unterminated string")
    return bytes(data[pos:end]), end + 1


def decode_save(raw: bytes) -> tuple[bytearray, int]:
    if len(raw) < 22 or raw[0] != len(VARIANT) or raw[1:1 + len(VARIANT)] != VARIANT:
        raise ValueError("SAVE26形式のHengbandセーブデータではありません")
    # length + unencoded variant + four chained version bytes + random XOR key
    start = 1 + len(VARIANT) + 4 + 1
    plain = bytearray(raw)
    for i in range(start, len(raw)):
        plain[i] = raw[i] ^ raw[i - 1]

    version_at = start
    header_at = 1 + len(VARIANT)
    previous = 0
    game_version = []
    for encoded in raw[header_at:header_at + 4]:
        game_version.append(encoded ^ previous)
        previous = encoded
    # Saves made before 3.0.2.3 contain twelve obsolete bytes here. The
    # current loader skips them in rd_version_info() before reading SAVE<n>.
    if tuple(game_version) < (3, 0, 2, 3):
        version_at += 12

    version = struct.unpack_from("<I", plain, version_at)[0]
    if version not in SUPPORTED_SAVE_VERSIONS:
        raise ValueError(
            f"ゲームバージョン {'.'.join(map(str, game_version))} のSAVE{version}には未対応です"
            "（対応形式: SAVE23～SAVE26）"
        )
    return plain, start


def verify_checksums(raw: bytes, plain: bytearray, start: int) -> None:
    if len(raw) < start + 8:
        raise ValueError("セーブデータが短すぎます")
    stored_v, stored_x = struct.unpack_from("<II", plain, len(plain) - 8)
    actual_v = sum(plain[start:-8]) & 0xFFFFFFFF
    actual_x = sum(raw[start:-4]) & 0xFFFFFFFF
    if stored_v != actual_v or stored_x != actual_x:
        raise ValueError(
            "チェックサムが一致しません: "
            f"decoded={stored_v:08x}/{actual_v:08x}, encoded={stored_x:08x}/{actual_x:08x}"
        )


def find_player(plain: bytearray, start: int, requested_name: str | None) -> tuple[int, bytes]:
    candidates: list[tuple[int, int, bytes]] = []
    # A player record starts with name, cause of death, and last message strings.
    # Scanning is used because all preceding lore/quest sections have variable size.
    for record in range(start, len(plain) - 400):
        try:
            name, p = cstring(plain, record, 64)
            if not name or len(name) > 31:
                continue
            _, p = cstring(plain, p, 1024)
            _, p = cstring(plain, p, 4096)
            p += QUICK_START_FIXED_PREFIX
            for _ in range(4):
                _, p = cstring(plain, p, 4096)
            p += 2  # unused random-quest count and quick_ok
            for _ in range(4):
                _, p = cstring(plain, p, 4096)

            # race,class,personality,sex,realm1,realm2,reserved,hit-die,expfact
            if plain[p + 6] != 0 or not 1 <= plain[p + 7] <= 100:
                continue
            expfact, age, height, weight = struct.unpack_from("<Hhhh", plain, p + 8)
            if not (10 <= expfact <= 1000 and 1 <= age <= 10000 and 10 <= height <= 1000 and 10 <= weight <= 5000):
                continue
            candidates.append((record, p + 12, name))
        except (ValueError, IndexError, struct.error):
            continue

    if requested_name:
        encodings = ("utf-8", "cp932", "euc_jp")
        candidates = [c for c in candidates if any(_decode(c[2], e) == requested_name for e in encodings)]

    # Starting the scan in the middle of "PLAYER" also produces "LAYER",
    # "AYER", etc. They all resolve to the same body offset; retain the
    # earliest (complete) name. False positives later in the record are also
    # eliminated by preferring the earliest structurally valid record.
    by_body: dict[int, tuple[int, int, bytes]] = {}
    for candidate in candidates:
        body = candidate[1]
        if body not in by_body or candidate[0] < by_body[body][0]:
            by_body[body] = candidate
    unique = sorted(by_body.values(), key=lambda c: c[0])
    if requested_name and len(unique) != 1:
        names = ", ".join(display_text(n) for _, _, n in unique) or "なし"
        raise ValueError(f"プレイヤー領域を一意に特定できませんでした（候補: {names}）。--nameを指定してください")
    if not unique:
        raise ValueError("プレイヤー領域を特定できませんでした")
    _, body, name = unique[0]
    return body, name


def _decode(value: bytes, encoding: str) -> str:
    try:
        return value.decode(encoding)
    except UnicodeDecodeError:
        return ""


def display_text(value: bytes) -> str:
    for encoding in ("utf-8", "cp932", "euc_jp"):
        text = _decode(value, encoding)
        if text and text.isprintable():
            return text
    return value.hex()


def find_social_class(plain: bytearray, body_at: int) -> int:
    """Locate sc using its reserved-field signature in wr_player()."""
    candidates = []
    end = min(len(plain) - 12, body_at + 20000)
    for reserved_at in range(body_at + 100, end):
        if plain[reserved_at:reserved_at + 8] != b"\0" * 8:
            continue
        social = struct.unpack_from("<h", plain, reserved_at + 8)[0]
        old_rest = struct.unpack_from("<h", plain, reserved_at + 10)[0]
        if 1 <= social <= 100 and old_rest == 0:
            candidates.append(reserved_at + 8)
    if not candidates:
        raise ValueError("社会的地位の保存位置を特定できませんでした")
    # sc is the first field with this signature. Similar zero-filled reserved
    # fields can occur later among timed effects and other player state.
    return candidates[0]


def encode_with_checksums(raw: bytes, plain: bytearray, start: int) -> bytes:
    struct.pack_into("<I", plain, len(plain) - 8, sum(plain[start:-8]) & 0xFFFFFFFF)
    out = bytearray(raw[:start]) + bytearray(len(raw) - start)
    previous = raw[start - 1]
    for i in range(start, len(raw) - 4):
        previous ^= plain[i]
        out[i] = previous
    xstamp = sum(out[start:-4]) & 0xFFFFFFFF
    struct.pack_into("<I", plain, len(plain) - 4, xstamp)
    for i in range(len(raw) - 4, len(raw)):
        previous ^= plain[i]
        out[i] = previous
    return bytes(out)


def main() -> int:
    parser = argparse.ArgumentParser(description="Hengband SAVE26の身長・体重を表示・変更します")
    parser.add_argument("savefile", type=Path)
    parser.add_argument("--name", help="候補が複数ある場合のキャラクター名")
    parser.add_argument("--age", type=int, metavar="YEARS", help="新しい年齢")
    parser.add_argument("--height", type=int, metavar="CM", help="新しい身長（cm）")
    parser.add_argument("--weight", type=int, metavar="KG", help="新しい体重（kg）")
    parser.add_argument("--social-class", type=int, metavar="NUMBER", help="新しい社会的地位（1～100）")
    parser.add_argument("--in-place", action="store_true", help="元ファイルを書換え（既定動作、.bakも作成）")
    parser.add_argument("--output", type=Path, help="書換えたセーブデータの出力先")
    args = parser.parse_args()

    try:
        raw = args.savefile.read_bytes()
        plain, start = decode_save(raw)
        verify_checksums(raw, plain, start)
        body_at, name = find_player(plain, start, args.name)
        social_at = find_social_class(plain, body_at)
        age = struct.unpack_from("<h", plain, body_at - 2)[0]
        height_in, weight_lb = struct.unpack_from("<hh", plain, body_at)
        social_class = struct.unpack_from("<h", plain, social_at)[0]
        print(f"名前: {display_text(name)}")
        print(f"年齢: {age}")
        print(f"身長: {height_in * 254 // 100} cm（内部値 {height_in} inch）")
        print(f"体重: {weight_lb * 4536 // 10000} kg（内部値 {weight_lb} lb）")
        print(f"社会的地位: {social_class}")

        if args.age is None and args.height is None and args.weight is None and args.social_class is None:
            return 0
        if args.in_place and args.output:
            parser.error("--in-placeと--outputは同時に指定できません")

        new_age = age if args.age is None else args.age
        new_height = height_in if args.height is None else round(args.height * 100 / 254)
        new_weight = weight_lb if args.weight is None else round(args.weight * 10000 / 4536)
        new_social = social_class if args.social_class is None else args.social_class
        if not (-32768 <= new_age <= 32767 and -32768 <= new_height <= 32767 and -32768 <= new_weight <= 32767):
            raise ValueError("指定値がセーブ形式の16ビット範囲外です")
        if not 1 <= new_age <= 10000:
            raise ValueError("年齢は1～10000で指定してください")
        if not 1 <= new_social <= 100:
            raise ValueError("社会的地位は1～100で指定してください")
        struct.pack_into("<h", plain, body_at - 2, new_age)
        struct.pack_into("<hh", plain, body_at, new_height, new_weight)
        struct.pack_into("<h", plain, social_at, new_social)
        result = encode_with_checksums(raw, plain, start)
        check_plain, check_start = decode_save(result)
        verify_checksums(result, check_plain, check_start)

        destination = args.output if args.output is not None else args.savefile
        if args.output is None:
            backup = args.savefile.with_name(args.savefile.name + ".bak")
            shutil.copy2(args.savefile, backup)
            print(f"バックアップ: {backup}")
        destination.write_bytes(result)
        print(f"書込み先: {destination}")
        print(
            f"変更後: 年齢 {new_age}、身長 {new_height * 254 // 100} cm、"
            f"体重 {new_weight * 4536 // 10000} kg、社会的地位 {new_social}"
        )
        return 0
    except (OSError, ValueError) as exc:
        print(f"エラー: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())