#!/usr/bin/env python3
"""feature_scan.py — Inventory the feature names used by TA map archives.

Walks a folder of .zip maps, reads every TNT's feature table and writes a
JSON mapping: {zip filename: [feature name, ...]}.

Usage:
    python feature_scan.py /path/to/maps results.json

Format notes and the online feature inventory: https://ta-archive.com/tools/py/
MIT licence.
"""
import json
import os
import re
import struct
import sys
import zipfile

SENTINELS = {0xFFFF, 0xFFFC, 0xFFFE}
# fallback for maps whose zips include the features directory instead
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
try:
    from hpi import Hpi, collect, extract_file  # same folder
except ImportError:
    Hpi = None


def feature_table(tnt):
    """Extract feature names from a TNT blob: (names, used_indices)."""
    if len(tnt) < 64:
        return [], []
    hdr = struct.unpack('<16I', tnt[:64])
    if hdr[0] != 0x2000:
        return [], []
    w, h = hdr[1], hdr[2]
    if not (0 < w < 500 and 0 < h < 500):
        return [], []
    n = w * h
    attrs = tnt[hdr[4]:hdr[4] + n * 4]
    if len(attrs) < n * 4:
        return [], []
    used = set()
    for i in range(n):
        rec = attrs[i * 4:i * 4 + 4]
        fi = rec[1] | (rec[2] << 8)
        if fi not in SENTINELS:
            used.add(fi)
    if not used or max(used) > 500:
        return [], []
    count = max(used) + 1
    # locate the table: records are 132 bytes = 4B prefix + 128B NUL-padded name,
    # sitting between the tile graphics and the 252x252 minimap block.
    for m in re.finditer(rb'[\x00][\x00]{3}[ -~][ -~]{2,60}\x00', tnt):
        ptr = m.start()
        cand, ok = [], True
        for idx in sorted(used)[:8]:
            rec = tnt[ptr + idx * 132:ptr + idx * 132 + 132]
            if len(rec) < 132:
                ok = False
                break
            nm = rec[4:132].split(b'\x00')[0]
            if not nm or not all(32 <= c < 127 for c in nm):
                ok = False
                break
            cand.append(nm.decode('latin-1'))
        if ok and cand:
            full = []
            for idx in range(count):
                rec = tnt[ptr + idx * 132:ptr + idx * 132 + 132]
                nm = rec[4:132].split(b'\x00')[0] if len(rec) >= 132 else b''
                full.append(nm.decode('latin-1') if nm else None)
            return full, sorted(used)
    return [], sorted(used)


def scan_zip(path):
    """All feature names referenced by the maps inside one zip."""
    feats = set()
    try:
        z = zipfile.ZipFile(path)
    except Exception:
        return feats
    for n in z.namelist():
        lower = n.lower()
        if lower.endswith('.tnt'):
            try:
                feats.update(x for x in feature_table(z.read(n))[0] if x)
            except Exception:
                pass
        elif lower.endswith(('.ufo', '.hpi', '.two', '.ccx')) and Hpi is not None:
            try:
                h = Hpi(z.read(n))
                for name, info in collect(h.directory()):
                    if name.lower().endswith('.tnt'):
                        feats.update(
                            x for x in feature_table(extract_file(h, info))[0] if x)
            except Exception:
                pass
    return feats


def main():
    if len(sys.argv) < 3:
        print(__doc__)
        sys.exit(1)
    src_dir, dst = sys.argv[1], sys.argv[2]
    results = {}
    zips = sorted(f for f in os.listdir(src_dir) if f.lower().endswith('.zip'))
    for i, zn in enumerate(zips):
        try:
            feats = scan_zip(os.path.join(src_dir, zn))
        except Exception:
            feats = set()
        results[zn] = sorted(feats)
        if i % 50 == 0:
            print(f'{i}/{len(zips)}', flush=True)
    json.dump(results, open(dst, 'w'), indent=1)
    names = set()
    for v in results.values():
        names.update(v)
    print(f'done: {len(results)} maps, {len(names)} distinct feature names')


if __name__ == '__main__':
    main()