#!/usr/bin/env python3
"""Extract autochess art from a locally installed Arknights client (optional, host-side).

The official CN client ships Unity AssetBundles compressed with a custom LZ4 variant ("LZ4AK");
aklz4.py registers a decoder for it. This script pulls the art the web sources lack:
  - map_autochess / map_autochesssand textures (the real board tile atlas)
  - every Sprite of the three autochess UI bundles (HUD frames, badges, banners, icons)
  - the 6 in-match emoticon themes of 盟约 (activity_table autoChessData.enabledEmoticonThemeIdList; only the
    *_battle sprites, keyed by display_meta_table picId: emoticon/<dir>/<picId>.png, see data/emotes.json)
  - the autochess guidebook pages, battle projectile sprites and a few token/skin Spine models missing upstream
  - for the official 3D board (DESIGN §15): the map theme's Material parameters (map/<theme>/materials.json; shader
    names resolved through the shaders/*.ab bundles), the background / device meshes as Wavefront OBJ
    (mesh/<bundle>/<mesh>.obj; UnityPy's exporter, X mirrored into a right-handed frame) and their GameObject
    placement (mesh/<bundle>/prefab.json)
  - the standard gate / objective effects of every battle map (arts/effects/[pack]map.ab → map/fx: the [opt]start_box /
    [opt]end_box / [opt]start_fly meshes, their merged additive textures, materials and prefab), the blower's texture
    (arts/maps/common/res.ab → map/common: TX_wind_device of the s_wind_device mesh) and the shared water / noise
    textures (arts/maps/effect.ab → map/water) used for deep-sea / mire / smog tiles
  - derived PBR maps for three.js (DERIVED): Unity stores the theme's normal map as two channels (BC5: RG = XY, B = 0)
    and metallic/gloss with the smoothness in A; the board renderer needs an RGB normal map (Z rebuilt) and a
    roughness map in G (1 − smoothness; metalness B = 0), written next to the source as <name>_rgb.png / _rough.png

Usage:
  python3 -m venv .venv && .venv/bin/pip install -r tools/local-extract/requirements.txt
  .venv/bin/python tools/local-extract/extract.py [--game <AB root>] [--out public/assets/local] [--only <subdir prefix>]
  python3 tools/local-extract/extract.py --print-jobs     (the job table as JSON; needs no dependencies)

Writes <out>/**.png|.skel|.atlas and data/local-assets.json (manifest of what was extracted). With --only, just the
jobs whose output subdir starts with one of the prefixes run, and their groups replace those of the existing manifest
(every other group is kept as is).
Everything is (c) Hypergryph; for private, non-commercial fan use only.
"""
import argparse
import json
import os
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

ROOT = Path(__file__).resolve().parents[2]
HOME = Path.home()
CANDIDATES = [
    HOME / 'Library/Application Support/CrossOver/Bottles/Arknights/drive_c/Program Files/Hypergryph Launcher/games/Arknights/Arknights_Data/StreamingAssets/AB/Windows',
    HOME / 'Library/Containers/com.hypergryph.arknights/Data/Documents/Bundles',
    # native Windows installs of the official launcher (default and x86 locations)
    Path('C:/Program Files/Hypergryph Launcher/games/Arknights/Arknights_Data/StreamingAssets/AB/Windows'),
    Path('C:/Program Files (x86)/Hypergryph Launcher/games/Arknights/Arknights_Data/StreamingAssets/AB/Windows'),
]

# In-match emote themes (display_meta_table emoticonData, scene AUTOCHESS_BATTLE), in the order of
# activity_table autoChessData.enabledEmoticonThemeIdList: (themeId, output dir under emoticon/). The dirs are the
# ones of shared/constants.js EMOTE_THEMES / data/emotes.json. Only the *_battle sprites are exported: the bundles also
# hold room / pick pics of other modes (and the autochess room pics are not in the local client at all).
EMOTE_THEMES = [
    ('emoticon_autochess_basic', 'basic'),
    ('emoticon_originium_slug', 'slug'),
    ('emoticon_autochess_basic_2', 'basic_2'),
    ('emoticon_foolsday_doctor', 'fooldoctor'),
    ('emoticon_foolsday_amiya', 'foolamiya'),
    ('emoticon_foolsday_wisdel', 'foolwisdel'),
]
BATTLE_EMOTE = r'^pic_.+_battle$'

# Meshes of the official board scene (DESIGN §15): (bundle, output name under mesh/).
MESH_BUNDLES = [
    ('arts/maps/map_autochess/bkg_mesh.ab', 'map_autochess_bkg'),
    ('arts/maps/common/meshes/s_background_common.ab', 's_background_common'),
    ('arts/maps/common/meshes/s_common_box_01.ab', 's_common_box_01'),
    ('arts/maps/common/meshes/s_wind_device.ab', 's_wind_device'),
]

# (bundle path relative to the AB root, output subdir, which object types to export[, name regex to keep])
JOBS = [
    ('arts/maps/map_autochess/res.ab', 'map/autochess', {'Texture2D', 'Material'}),
    ('arts/maps/map_autochesssand/res.ab', 'map/autochesssand', {'Texture2D', 'Material'}),
    ('ui/autochess/[uc]autochesscommon.ab', 'ui/common', {'Sprite'}),
    ('ui/autochess/[uc]autochessbattle.ab', 'ui/battle', {'Sprite', 'TextAsset'}),
    ('ui/autochess/[uc]autochessouter.ab', 'ui/outer', {'Sprite'}),
    *[(f'ui/emoticon/theme/[uc]{theme}.ab', f'emoticon/{sub}', {'Sprite'}, BATTLE_EMOTE) for theme, sub in EMOTE_THEMES],
    ('arts/guidebookpages/[pack]autochess.ab', 'guide', {'Sprite', 'Texture2D'}),
    ('battle/prefabs/[uc]projectiles.ab', 'projectiles', {'Sprite', 'Texture2D'}),
    # official module (uniequip) type icons, keyed by lower-case type name (e.g. 'mar-x')
    ('spritepack/ui_equip_type_hub_h2_0.ab', 'module', {'Sprite'}),
    ('skinpack/token_10039_ulpia_block.ab', 'spine/token_10039_ulpia_block', {'TextAsset', 'Texture2D'}),
    *[(rel, f'mesh/{sub}', {'Mesh', 'GameObject', 'Material', 'Texture2D'}) for rel, sub in MESH_BUNDLES],
    ('arts/effects/[pack]map.ab', 'map/fx', {'Mesh', 'GameObject', 'Material', 'Texture2D'}),
    ('arts/maps/common/res.ab', 'map/common', {'Material', 'Texture2D'}, r'^(TX|MT)_wind_device$'),
    ('arts/maps/effect.ab', 'map/water', {'Texture2D'}, r'water_normal|Caustics|WaterNoise|noise_clouds|Water_Foam|SmoothWaves'),
]

# Derived three.js maps: (output subdir, source texture, kind, output name). 'normal_rg' rebuilds Z of a two-channel
# (BC5) normal map into an RGB tangent-space map; 'rough_from_gloss' turns Unity's metallic/gloss map (smoothness in
# A) into a roughness (G) / metalness (B) map. Written after the job that exported the source texture.
DERIVED = [
    ('map/autochess', 'TX_autochessi_N', 'normal_rg', 'TX_autochessi_N_rgb'),
    ('map/autochess', 'TX_autochessi_M', 'rough_from_gloss', 'TX_autochessi_M_rough'),
]

# Shader bundles loaded beside every job that exports Materials, only so that the materials' shader references
# (external CABs) resolve to a name in materials.json (e.g. Torappu/Scene/StandardDirectional); nothing is exported
# from them.
SHADER_DEPS = 'shaders/*.ab'
# the root shader bundle ([uc]shaders.ab) holds the Torappu particle shaders of the map effects (map/fx)
SHADER_DEPS_ROOT = '[uc]shaders.ab'

SAFE = re.compile(r'[^A-Za-z0-9_.\-\[\]]+')


def safe_name(name):
    name = SAFE.sub('_', name).strip('._') or 'unnamed'
    return name[:120]


def job_parts(job):
    """(rel, sub, kinds, keep) of a JOBS entry; keep is a compiled name filter or None."""
    rel, sub, kinds = job[:3]
    keep = re.compile(job[3]) if len(job) > 3 and job[3] else None
    return rel, sub, kinds, keep


def select_jobs(only):
    """The jobs whose output subdir starts with one of the `only` prefixes (all jobs when `only` is empty)."""
    if not only:
        return list(JOBS)
    return [j for j in JOBS if any(j[1] == p.rstrip('/') or j[1].startswith(p.rstrip('/') + '/') for p in only)]


def _num(v, nd=6):
    """JSON-safe rounded float (NaN / inf → 0)."""
    try:
        f = float(v)
    except (TypeError, ValueError):
        return 0
    return round(f, nd) + 0.0 if f == f and abs(f) != float('inf') else 0  # + 0.0: no '-0.0'


def vec(v, keys='xyz'):
    """[x, y, z(, w)] of a UnityPy vector / quaternion (missing components → 0)."""
    return [_num(getattr(v, k, 0)) for k in keys]


def _pptr_name(pptr):
    """Name of the object a PPtr points to, or None (null pointer or an unloaded external bundle)."""
    try:
        if not getattr(pptr, 'm_PathID', 0):
            return None
        return getattr(pptr.read(), 'm_Name', None) or None
    except Exception:  # external dependency not loaded, unsupported type, …
        return None


def material_info(mat):
    """JSON summary of a Material: shader, keywords, textures (by name, with tiling), floats and colours."""
    shader = None
    try:
        sh = mat.m_Shader.read()
        shader = getattr(getattr(sh, 'm_ParsedForm', None), 'm_Name', None) or getattr(sh, 'm_Name', None)
    except Exception:
        shader = None
    props = getattr(mat, 'm_SavedProperties', None)
    textures, floats, colors = {}, {}, {}
    for k, env in (getattr(props, 'm_TexEnvs', None) or []):
        name = _pptr_name(env.m_Texture)
        if name:
            textures[k] = {'texture': name, 'scale': vec(env.m_Scale, 'xy'), 'offset': vec(env.m_Offset, 'xy')}
    for k, v in (getattr(props, 'm_Floats', None) or []):
        floats[k] = _num(v)
    for k, c in (getattr(props, 'm_Colors', None) or []):
        colors[k] = vec(c, 'rgba')
    kw = getattr(mat, 'm_ValidKeywords', None) or getattr(mat, 'm_ShaderKeywords', None) or []
    if isinstance(kw, str):
        kw = kw.split()
    return {'shader': shader, 'keywords': sorted(kw), 'textures': dict(sorted(textures.items())),
            'floats': dict(sorted(floats.items())), 'colors': dict(sorted(colors.items()))}


def prefab_node(go, mesh_keys=None):
    """JSON summary of a GameObject: local transform (Unity space), parent, mesh and material names. `mesh_keys`
    ({mesh path id: manifest key}) names the exported OBJ when two meshes of a bundle share a name."""
    node = {'name': getattr(go, 'm_Name', '') or '', 'parent': None, 'pos': [0, 0, 0], 'rot': [0, 0, 0, 1],
            'scale': [1, 1, 1], 'mesh': None, 'materials': []}
    for c in getattr(go, 'm_Component', None) or []:
        pptr = getattr(c, 'component', c)
        try:
            comp = pptr.read()
        except Exception:
            continue
        kind = type(comp).__name__
        if kind in ('Transform', 'RectTransform'):
            node['pos'] = vec(comp.m_LocalPosition)
            node['rot'] = vec(comp.m_LocalRotation, 'xyzw')
            node['scale'] = vec(comp.m_LocalScale)
            try:
                father = comp.m_Father.read() if getattr(comp.m_Father, 'm_PathID', 0) else None
                node['parent'] = father.m_GameObject.read().m_Name if father else None
            except Exception:
                node['parent'] = None
        elif kind == 'MeshFilter':
            node['mesh'] = _pptr_name(comp.m_Mesh)
            pid = getattr(comp.m_Mesh, 'm_PathID', 0)
            if mesh_keys and pid in mesh_keys and not getattr(comp.m_Mesh, 'm_FileID', 0):
                node['mesh'] = mesh_keys[pid]
        elif kind in ('MeshRenderer', 'SkinnedMeshRenderer'):
            node['materials'] = [_pptr_name(m) for m in getattr(comp, 'm_Materials', None) or []]
    return node


def dep_bundles(ab_root, kinds):
    """Bundles to load next to a job so its references resolve: the shader bundles when it exports Materials."""
    if 'Material' not in kinds:
        return []
    root = Path(ab_root)
    extra = [root / SHADER_DEPS_ROOT] if (root / SHADER_DEPS_ROOT).is_file() else []
    return sorted(p for p in root.glob(SHADER_DEPS) if p.is_file()) + extra


_NORMAL_Z = None


def _normal_z_table():
    """Z byte of a unit normal for every (x byte, y byte) pair: 128 + 127·sqrt(max(0, 1 − x² − y²))."""
    global _NORMAL_Z
    if _NORMAL_Z is None:
        t = bytearray(65536)
        for xb in range(256):
            x = xb / 127.5 - 1.0
            for yb in range(256):
                y = yb / 127.5 - 1.0
                z = max(0.0, 1.0 - x * x - y * y) ** 0.5
                t[(xb << 8) | yb] = min(255, int(round(127.5 + 127.5 * z)))
        _NORMAL_Z = bytes(t)
    return _NORMAL_Z


def derive_normal_rg(img):
    """RGB tangent-space normal map from a two-channel (BC5: R = X, G = Y, B unused) one: Z rebuilt per pixel."""
    from PIL import Image
    rgb = img.convert('RGB')
    r, g, _ = rgb.split()
    table = _normal_z_table()
    z = bytes(table[(x << 8) | y] for x, y in zip(r.tobytes(), g.tobytes()))
    return Image.merge('RGB', (r, g, Image.frombytes('L', rgb.size, z)))


def derive_rough_from_gloss(img):
    """three.js roughness / metalness map (G = 1 − smoothness from Unity's A, B = metalness 0, R = 1) from a
    Unity metallic/gloss map."""
    from PIL import Image
    rgba = img.convert('RGBA')
    rough = rgba.getchannel('A').point(lambda v: 255 - v)
    return Image.merge('RGB', (Image.new('L', rgba.size, 255), rough, Image.new('L', rgba.size, 0)))


DERIVERS = {'normal_rg': derive_normal_rg, 'rough_from_gloss': derive_rough_from_gloss}


def run_derived(out_root, sub, manifest, log):
    """Write the DERIVED maps of output subdir `sub` (from its exported PNGs) and record them in the manifest."""
    from PIL import Image
    n = 0
    for dsub, src, kind, name in DERIVED:
        if dsub != sub:
            continue
        path = Path(out_root) / sub / f'{src}.png'
        if not path.exists():
            log(f'  derived {name}: source {src} missing')
            continue
        try:
            img = DERIVERS[kind](Image.open(path))
            img.save(Path(out_root) / sub / f'{name}.png')
        except Exception as e:  # a broken source only drops the derived map (the renderer falls back)
            log(f'  warn derived {name}: {e}')
            continue
        manifest.setdefault(sub, {})[name] = {'path': f'/assets/local/{sub}/{name}.png', 'w': img.width, 'h': img.height,
                                              'kind': 'Derived', 'from': src, 'derive': kind}
        n += 1
    return n


def merge_manifest(old_groups, new_groups, ran_subs):
    """Groups of the previous manifest minus the subdirs that were re-extracted, plus the new ones (sorted keys)."""
    out = {g: v for g, v in (old_groups or {}).items() if g not in ran_subs}
    out.update(new_groups)
    return {g: dict(sorted(out[g].items())) for g in sorted(out)}


def export_bundle(ab_root, job, out_root, manifest, log):
    import aklz4  # noqa: F401  (registers the LZ4AK decoder)
    import UnityPy
    rel, sub, kinds, keep = job_parts(job)
    src = ab_root / rel
    if not src.exists():
        log(f'skip (missing) {rel}')
        return 0
    try:
        env = UnityPy.load(str(src))
    except Exception as e:  # corrupt or unsupported bundle: report and continue
        log(f'FAIL load {rel}: {e}')
        return 0
    objects = list(env.objects)  # the job's own objects, listed before any dependency joins the environment
    for dep in dep_bundles(ab_root, kinds):
        try:
            env.load_file(str(dep))
        except Exception as e:  # a missing shader only leaves that material's shader name null
            log(f'  warn dependency {dep.name}: {e}')
    out_dir = out_root / sub
    out_dir.mkdir(parents=True, exist_ok=True)
    seen, n = set(), 0
    materials, prefab, gos, mesh_keys = {}, [], [], {}
    for obj in objects:
        t = obj.type.name
        if t not in kinds:
            continue
        try:
            data = obj.read()
            raw_name = getattr(data, 'm_Name', '') or ''
            if keep is not None and not keep.search(raw_name):
                continue
            name = safe_name(raw_name or f'obj_{obj.path_id}')
            if t in ('Sprite', 'Texture2D'):
                img = data.image
                if img is None or img.width < 2 or img.height < 2:
                    continue
                fname = name + '.png'
                if fname in seen:  # same name twice: a Sprite (true aspect) beats its padded Texture2D
                    if t == 'Texture2D' or manifest.get(sub, {}).get(name, {}).get('kind') == 'Sprite':
                        continue
                img.save(out_dir / fname)
                seen.add(fname)
                manifest.setdefault(sub, {})[name] = {
                    'path': f'/assets/local/{sub}/{fname}', 'w': img.width, 'h': img.height, 'kind': t}
                n += 1
            elif t == 'TextAsset':
                raw = data.m_Script
                blob = raw.encode('utf-8', 'surrogateescape') if isinstance(raw, str) else bytes(raw)
                if name.endswith('.atlas') or name.endswith('.skel'):
                    fname = name
                elif blob[:1] in (b'\n', b'') or b'size:' in blob[:200]:
                    fname = name + '.atlas'
                else:
                    fname = name + '.skel'
                (out_dir / fname).write_bytes(blob)
                manifest.setdefault(sub, {})[fname] = {'path': f'/assets/local/{sub}/{fname}', 'kind': t}
                n += 1
            elif t == 'Mesh':
                text = data.export()
                if not text or 'v ' not in text:
                    continue
                fname = name + '.obj'
                if fname in seen:
                    name, fname = f'{name}_{obj.path_id}', f'{name}_{obj.path_id}.obj'
                (out_dir / fname).write_text(text, encoding='utf-8')
                seen.add(fname)
                verts = sum(1 for line in text.splitlines() if line.startswith('v '))
                manifest.setdefault(sub, {})[name] = {'path': f'/assets/local/{sub}/{fname}', 'kind': 'Mesh', 'verts': verts}
                mesh_keys[obj.path_id] = name
                n += 1
            elif t == 'Material':
                materials[raw_name or name] = material_info(data)
            elif t == 'GameObject':
                gos.append(data)  # summarised after every mesh got its manifest key
        except Exception as e:
            log(f'  warn {rel} #{obj.path_id}: {e}')
    for go in gos:
        try:
            prefab.append(prefab_node(go, mesh_keys))
        except Exception as e:
            log(f'  warn {rel} GameObject: {e}')
    for key, fname, payload in (('materials', 'materials.json', materials), ('prefab', 'prefab.json', prefab)):
        if not payload:
            continue
        if isinstance(payload, list):
            payload = sorted(payload, key=lambda nd: (nd['parent'] or '', nd['name']))
        else:
            payload = dict(sorted(payload.items()))
        (out_dir / fname).write_text(json.dumps(payload, ensure_ascii=False, separators=(',', ':')), encoding='utf-8')
        manifest.setdefault(sub, {})[key] = {'path': f'/assets/local/{sub}/{fname}', 'kind': key.capitalize(), 'count': len(payload)}
        n += 1
    n += run_derived(out_root, sub, manifest, log)
    log(f'{rel}: {n} files -> {sub}')
    return n


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('--game', help='AssetBundle root (…/StreamingAssets/AB/Windows or …/Documents/Bundles)')
    ap.add_argument('--out', default=str(ROOT / 'public/assets/local'))
    ap.add_argument('--manifest', default=str(ROOT / 'data/local-assets.json'))
    ap.add_argument('--only', action='append', default=[], metavar='SUBDIR',
                    help='only run the jobs whose output subdir starts with this prefix (repeatable), e.g. emoticon')
    ap.add_argument('--print-jobs', action='store_true', help='print the job table as JSON and exit')
    args = ap.parse_args()

    if args.print_jobs:
        jobs = [{'bundle': rel, 'sub': sub, 'kinds': sorted(kinds), 'keep': keep.pattern if keep else None}
                for rel, sub, kinds, keep in map(job_parts, JOBS)]
        derived = [{'sub': sub, 'from': src, 'derive': kind, 'name': name} for sub, src, kind, name in DERIVED]
        print(json.dumps({'emoteThemes': [{'themeId': t, 'dir': d} for t, d in EMOTE_THEMES], 'jobs': jobs,
                          'derived': derived}, ensure_ascii=False))
        return 0

    ab_root = Path(args.game) if args.game else next((p for p in CANDIDATES if p.exists()), None)
    if not ab_root or not ab_root.exists():
        print('No Arknights install found. Pass --game <AssetBundle root>.', file=sys.stderr)
        return 2
    jobs = select_jobs(args.only)
    if not jobs:
        print(f'--only {args.only}: no job matches', file=sys.stderr)
        return 2
    out_root = Path(args.out)
    out_root.mkdir(parents=True, exist_ok=True)
    manifest, total = {}, 0
    print(f'AB root: {ab_root}')
    for job in jobs:
        total += export_bundle(ab_root, job, out_root, manifest, print)
    old = {}
    if args.only and Path(args.manifest).exists():
        try:
            old = json.loads(Path(args.manifest).read_text(encoding='utf-8')).get('groups') or {}
        except (OSError, ValueError) as e:
            print(f'warn: previous manifest unreadable ({e}); writing only the re-extracted groups', file=sys.stderr)
    groups = merge_manifest(old, manifest, set(manifest))  # a missing bundle keeps its previous group
    count = sum(len(v) for v in groups.values())
    doc = {'version': 1, 'source': 'local-client', 'count': count, 'groups': groups}
    Path(args.manifest).write_text(json.dumps(doc, ensure_ascii=False, separators=(',', ':')), encoding='utf-8')
    print(f'done: {total} files extracted, manifest {args.manifest} ({count} entries)')
    return 0 if total else 1


if __name__ == '__main__':
    sys.exit(main())