"""Lighter copy of one library asset: fewer triangles, same look, saved as a new DRAFT version of the same id (games keep
the active version until someone presses 'Use in games').
usage: python3 tools/lighter.py <asset id> [--target <triangles>] [--max-change <pct>]   (no --target = auto)
       python3 tools/lighter.py --activate <asset id> <version>
Auto tries growing error budgets with glTF-Transform simplify (meshoptimizer), renders every try against the original on
the gex44 GPU from 4 sides (tools/gpu/lighter_compare.sh) and keeps the lightest whose worst side changes at most
--max-change % of the model's pixels and no small spot changes by more than MAX_SPOT in brightness. A copy keeps the original's geometry compression (Draco / meshopt) and textures.
Prints the new version number. Exit 3 = no copy worth making (the reason is the last stderr line)."""
import copy, hashlib, json, os, shutil, subprocess, sys, time
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, 'lib'))
import glb, card as cardlib  # noqa: E402

GT = ['npx', '-y', '@gltf-transform/cli@4.5.1']
ERRORS = [0.0005, 0.001, 0.002, 0.003, 0.004, 0.006, 0.008, 0.016, 0.03]  # auto ladder: meshoptimizer's error budget (share of the model's size)
MAX_CHANGE = 5.0       # % of the model's pixels allowed to change on its worst side. Trawler, 10 Oct: 3.7 % and 6 % look the
                       # same as the original, 12.5 % shows bent rails and a warped hull stripe
MIN_SAVING = 0.15      # a copy that saves less than this is not worth a version
MAX_SPOT = 50          # biggest brightness change (0-255) of an 8 x 8 px spot. Area % misses small holes: the Breton house
                       # at error 0.03 broke a side window yet changed only 1.7 % of its pixels (spot 113; fine tries 24-41)
MIN_TRIS = 5000        # below this a model is already light; on simple shapes the pixel check also misses profile changes
NO_RIG = 'rigged and animated people keep their mesh (skin weights, morph targets): the Cast lane makes their light versions'


def gt(*args):
    r = subprocess.run(GT + list(args), capture_output=True, text=True, timeout=600)
    if r.returncode:
        raise SystemExit(f'glTF-Transform failed: {(r.stderr or r.stdout).strip().splitlines()[-1][:300]}')


def make(src, out, ratio, error, codec):
    raw = out + '.raw.glb'
    gt('simplify', src, raw, '--ratio', f'{ratio:.4f}', '--error', str(error))
    if codec in ('draco', 'meshopt'):  # simplify decodes the geometry; put the original's compression back
        gt(codec, raw, out)
        os.remove(raw)
    else:
        os.replace(raw, out)
    return glb.measure(out)


def compare(row, jd, files):
    spec = f'{jd}/spec.json'
    json.dump({'row': row, 'files': [os.path.basename(f) for f in files]}, open(spec, 'w'))
    r = subprocess.run([f'{ROOT}/tools/gpu/lighter_compare.sh', spec, f'{jd}/compare', *files], capture_output=True, text=True, timeout=900)
    res = f'{jd}/compare/result.json'
    if r.returncode or not os.path.exists(res):
        raise SystemExit(f'GPU check failed: {(r.stderr or r.stdout).strip().splitlines()[-1:] or "no output"}')
    out = json.load(open(res))
    if not out.get('ok'):
        raise SystemExit(f"GPU check failed: {out.get('error')}")
    return out['results']


def sheet(jd, pick):
    """compare.jpg: top row the original from 4 sides, bottom row the copy that was kept."""
    from PIL import Image
    rows = [[Image.open(f'{jd}/compare/{i}_{k}.jpg') for k in range(4)] for i in (0, pick)]
    im = Image.new('RGB', (4 * 512, 2 * 512), 'white')
    for y, r in enumerate(rows):
        for x, t in enumerate(r):
            im.paste(t, (x * 512, y * 512))
    im.resize((1600, 800)).save(f'{jd}/compare.jpg', quality=86)


def lighter(aid, target=None, max_change=MAX_CHANGE):
    allc = cardlib.load_all(ROOT)
    if aid not in allc:
        raise SystemExit(f'unknown asset {aid}')
    src = cardlib.active(allc[aid]); ver = allc[aid][-1]['version'] + 1
    cls, slug = aid.split('/')
    if cls.split('.')[0] in ('character', 'motion') or cls.startswith('part.character') or src.get('rig'):
        raise SystemExit(NO_RIG)
    tris0, codec = src['render']['tris'], src['render']['codec']
    if tris0 < MIN_TRIS:
        print(f'already light: {tris0:,} triangles (lighter copies start at {MIN_TRIS:,})', file=sys.stderr)
        sys.exit(3)
    if target is not None and not 100 <= target < tris0:
        raise SystemExit(f'target must be between 100 and {tris0 - 1:,} triangles (the model has {tris0:,})')
    jd = f'{ROOT}/jobs/lighter-{slug}-v{ver}'
    if os.path.isdir(jd):
        shutil.rmtree(jd)
    os.makedirs(jd)
    sglb = src['files']['glb'].split(':')[1]
    source = f'{jd}/source.glb'
    shutil.copyfile(f'{ROOT}/store/blobs/{sglb}.glb', source)
    tries = []  # (file, measure, ratio, error)
    if target is not None:
        f = f'{jd}/target.glb'  # error unconstrained: as close to the target as the texture seams allow
        tries.append((f, make(source, f, target / tris0, 1, codec), target / tris0, 1))
    else:
        seen = set()
        for e in ERRORS:
            f = f'{jd}/e{e}.glb'
            m = make(source, f, 0.05, e, codec)
            if m['tris'] in seen or m['tris'] >= tris0 * (1 - MIN_SAVING / 2):
                os.remove(f); continue  # same result as a smaller budget, or barely lighter: not worth a GPU check
            seen.add(m['tris']); tries.append((f, m, 0.05, e))
        if not tries:
            done(jd)
            print(f'already lean: even the largest error budget keeps {tris0:,} triangles within a few %', file=sys.stderr)
            sys.exit(3)
    row = {'id': aid, 'kit': src['provenance'].get('kit'), 'scale': src['size'].get('scale_to_class', 1)}
    res = compare(row, jd, [source] + [t[0] for t in tries])
    worst = [r['worst'] for r in res[1:]]
    spot = [r['spot'] for r in res[1:]]
    if target is not None:
        pick = 0
    else:
        ok = [i for i, (w, sp) in enumerate(zip(worst, spot)) if w <= max_change and sp <= MAX_SPOT]
        if not ok:
            done(jd)
            print(f'already lean: the smallest cut ({tries[0][1]["tris"]:,} triangles) already changes {worst[0]:.1f} % '
                  f'of the model on its worst side (limit {max_change:g} %) or a spot by {spot[0]}/255 (limit {MAX_SPOT})', file=sys.stderr)
            sys.exit(3)
        pick = min(ok, key=lambda i: tries[i][1]['tris'])
    f, m, ratio, err = tries[pick]
    saving = 1 - m['tris'] / tris0
    if saving < MIN_SAVING:
        done(jd)
        print(f'already lean: the copy that keeps the look saves only {saving * 100:.0f} % of the triangles', file=sys.stderr)
        sys.exit(3)
    h = hashlib.sha256(open(f, 'rb').read()).hexdigest()
    shutil.copyfile(f, f'{ROOT}/store/blobs/{h}.glb')
    sheet(jd, pick + 1)
    c = copy.deepcopy(src)
    c.pop('active', None)
    c['version'] = ver
    c['status'] = 'draft'
    c['files'] = {**c['files'], 'glb': f'sha256:{h}'}
    c['render'] = {'tris': m['tris'], 'verts': m['verts'], 'materials': m['materials'], 'draw_calls': m['draw_calls'],
                   'textures': m['textures'], 'gpu_mb_decoded': m['gpu_mb_decoded'], 'bytes': m['bytes'], 'codec': m['codec'],
                   'extensions': m['extensions']}
    c['techniques'] = sorted(set(c.get('techniques') or []) | {'decimated'})
    sides = res[pick + 1]['sides']
    c['lighter'] = {'from': src['version'], 'tris_from': tris0, 'saving_pct': round(saving * 100, 1),
                    'change_pct': res[pick + 1]['worst'], 'sides_pct': [s['pct'] for s in sides], 'max_change_pct': max_change,
                    'spot': res[pick + 1]['spot'], 'max_spot': MAX_SPOT,
                    'mode': 'target' if target is not None else 'auto', 'target': target, 'ratio': round(ratio, 4), 'error': err,
                    'tried': [{'tris': t[1]['tris'], 'error': t[3], 'change_pct': w, 'spot': sp} for t, w, sp in zip(tries, worst, spot)],
                    'method': 'glTF-Transform 4.5.1 simplify (meshoptimizer), original compression re-applied',
                    'job': os.path.basename(jd), 'made': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}
    c['quality'] = {**c['quality'], 'gates': f"lighter copy of v{src['version']}: {m['tris']:,} of {tris0:,} triangles "
                    f"(-{saving * 100:.0f} %), worst side {res[pick + 1]['worst']:.1f} % of the model's pixels changed, "
                    f"worst spot {res[pick + 1]['spot']}/255"}
    c['release'] = cardlib.release_hash(c)
    p = cardlib.write(ROOT, c)
    done(jd)
    print(p, f"{m['tris']} tris (-{saving * 100:.0f} %), worst side {res[pick + 1]['worst']} %, spot {res[pick + 1]['spot']}", file=sys.stderr)
    print(ver)


def done(jd):
    """Keep the evidence (compare/result.json, compare.jpg = original over the kept copy), drop the tries and their
    renders: the kept copy lives in store/blobs."""
    for n in os.listdir(jd):
        if n.endswith('.glb') or n == 'spec.json':
            os.remove(f'{jd}/{n}')
    if os.path.isdir(f'{jd}/compare'):
        for n in os.listdir(f'{jd}/compare'):
            if n.endswith('.jpg'):
                os.remove(f'{jd}/compare/{n}')


def activate(aid, version):
    """Make <version> the one games get. A draft lighter copy takes its source's status (approved, legacy...)."""
    allc = cardlib.load_all(ROOT)
    if aid not in allc:
        raise SystemExit(f'unknown asset {aid}')
    vs = {c['version']: c for c in allc[aid]}
    if version not in vs:
        raise SystemExit(f'{aid} has no v{version}')
    t = vs[version]
    if t['status'] == 'deprecated':
        raise SystemExit(f'v{version} was withdrawn: it cannot go to games')
    if t['status'] == 'draft':
        t['status'] = vs[t['lighter']['from']]['status'] if t.get('lighter') and t['lighter']['from'] in vs else 'candidate'
    for c in allc[aid]:
        if c is t:
            c['active'] = True
        elif c.pop('active', None) is None:
            continue
        cardlib.write(ROOT, c)
    cardlib.write(ROOT, t)
    print(version)


if __name__ == '__main__':
    a = sys.argv[1:]
    if a and a[0] == '--activate':
        activate(a[1], int(a[2]))
    else:
        tg = int(a[a.index('--target') + 1]) if '--target' in a else None
        mc = float(a[a.index('--max-change') + 1]) if '--max-change' in a else MAX_CHANGE
        lighter(a[0], tg, mc)