#!/usr/bin/env python3
"""Harvest PSP firmware from every AMD Windows driver package for the RX 7900 XTX.

One package at a time, because disk is the binding constraint: each `.exe` is 0.6-1.5 GB
and a full NSIS extraction is ~3x that. So the loop is
    download -> extract ONLY amdkmdag.sys -> carve -> keep new blobs -> delete everything

Scoping note: firmware is extracted from `amdkmdag.sys` only. In 22.12.1 that file held
639 of the ~700 `$PS1` hits in the whole package; the rest were audio-driver and VCN
blobs, not PSP modules. So this is a deliberate, documented narrowing, not a claim that
nothing else carries firmware.

Resumable: every finished version is recorded in `state.json`, so re-running skips work
already done. Dedup is by body sha256 against the existing corpus (`psp_out/versions`,
`versions_lfw`, `versions_win`) plus everything harvested so far, so only genuinely new
builds are written.

usage: fw/harvest_drivers.py [--limit N] [--scratch DIR] [--out DIR]
"""
import argparse
import glob
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import time
import urllib.request

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from carve_driver import carve                                    # noqa: E402

PAGE = ("https://www.amd.com/en/support/downloads/previous-drivers.html/graphics/"
        "radeon-rx/radeon-rx-7000-series/amd-radeon-rx-7900-xtx.html")
UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) Chrome/120.0"
MIN_FREE_GB = 6


def free_gb(path):
    st = os.statvfs(path)
    return st.f_bavail * st.f_frsize / 2**30


def driver_list():
    req = urllib.request.Request(PAGE, headers={"User-Agent": UA})
    html = urllib.request.urlopen(req, timeout=60).read().decode("utf-8", "replace")
    urls = sorted(set(re.findall(r'https?://drivers\.amd\.com/[^"\'<> ]+\.exe', html)))
    out = []
    for u in urls:
        m = re.search(r'edition-([0-9]+\.[0-9]+(?:\.[0-9]+)?)', u)
        if not m:
            continue                                  # skip the RGB-LED utility etc.
        out.append((m.group(1), u))
    # oldest first
    def key(v):
        return [int(x) for x in v[0].split('.')]
    return sorted({v: u for v, u in out}.items(), key=key)


def corpus_hashes(repo):
    h = {}
    for d in ('psp_out/versions', 'psp_out/versions_lfw', 'psp_out/versions_win',
              'psp_out/versions_win/containers'):
        for p in glob.glob(os.path.join(repo, d, '*.bin')):
            try:
                h[hashlib.sha256(open(p, 'rb').read()).hexdigest()] = os.path.basename(p)
            except OSError:
                pass
    return h


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--scratch', default='/home/user/scratch/amd-drv/harvest')
    ap.add_argument('--out', default=None, help='default: <repo>/psp_out/versions_win')
    ap.add_argument('--limit', type=int, default=0)
    ap.add_argument('--repo', default=os.path.dirname(os.path.dirname(
        os.path.abspath(__file__))))
    args = ap.parse_args()

    out = args.out or os.path.join(args.repo, 'psp_out', 'versions_win')
    os.makedirs(args.scratch, exist_ok=True)
    os.makedirs(os.path.join(out, 'containers'), exist_ok=True)

    state_p = os.path.join(args.scratch, 'state.json')
    state = json.load(open(state_p)) if os.path.exists(state_p) else {}
    inv_p = os.path.join(args.repo, 'fw', 'windows_inventory.jsonl')

    known = corpus_hashes(args.repo)
    print('corpus: %d known bodies' % len(known), flush=True)

    drivers = driver_list()
    print('drivers: %d versions, %s .. %s'
          % (len(drivers), drivers[0][0], drivers[-1][0]), flush=True)

    done = 0
    for version, url in drivers:
        # retry anything that previously errored; skip only real successes
        if version in state and 'modules' in state[version]:
            continue
        if args.limit and done >= args.limit:
            break
        if free_gb(args.scratch) < MIN_FREE_GB:
            print('STOP: only %.1f GB free' % free_gb(args.scratch), flush=True)
            break
        exe = os.path.join(args.scratch, '%s.exe' % version)
        xdir = os.path.join(args.scratch, 'x_%s' % version)
        t0 = time.time()
        try:
            r = subprocess.run(['curl', '-sS', '-f', '-A', UA, '-e',
                                'https://www.amd.com/', '-o', exe, url],
                               capture_output=True, timeout=3600)
            if r.returncode != 0 or not os.path.exists(exe):
                state[version] = {'error': 'download rc=%d' % r.returncode}
                print('%-9s DOWNLOAD FAILED' % version, flush=True)
                continue
            # Two packaging formats, discovered the hard way: 22.x/23.0-23.5 are
            # NSIS self-extractors, while 23.7.1 onward are PE32+ 7-Zip SFX
            # (7z archive embedded at ~0x1b3e00).  `-tnsis` silently yields
            # nothing on the latter, which is why a first run reported
            # "NO amdkmdag.sys" for every version from 23.7.1.  Try both.
            sysfiles = []
            for fmt in ('-tnsis', '-t7z'):
                subprocess.run(['7z', 'x', '-y', fmt, '-o' + xdir, exe,
                                '-r', 'amdkmdag.sys'],
                               capture_output=True, timeout=3600)
                sysfiles = glob.glob(os.path.join(xdir, '**', 'amdkmdag.sys'),
                                     recursive=True)
                if sysfiles:
                    break
            if not sysfiles:
                state[version] = {'error': 'no amdkmdag.sys'}
                print('%-9s NO amdkmdag.sys' % version, flush=True)
                continue
            recs = []
            for sf in sysfiles:
                with open(sf, 'rb') as fh:
                    data = fh.read()
                recs.extend([r for r in carve(data, os.path.basename(sf))
                             if r['verdict'] == 'module'])
            new = 0
            seen = set()
            with open(inv_p, 'a') as inv:
                for r in recs:
                    if r['sha256'] in seen:
                        continue
                    seen.add(r['sha256'])
                    isnew = r['sha256'] not in known
                    rec = {k: v for k, v in r.items() if not k.startswith('_')}
                    rec['driver'] = version
                    rec['new'] = isnew
                    inv.write(json.dumps(rec) + '\n')
                    if isnew:
                        known[r['sha256']] = version
                        tag = (r['ta_version'] or r['version']).replace('.', '-')
                        nm = '%s_%s_%s.bin' % (r['guess'], tag, r['sha256'][:8])
                        with open(os.path.join(out, nm), 'wb') as fh:
                            fh.write(r['_body'])
                        with open(os.path.join(out, 'containers', nm), 'wb') as fh:
                            fh.write(r['_container'])
                        new += 1
            state[version] = {'modules': len(seen), 'new': new,
                              'secs': round(time.time() - t0, 1)}
            print('%-9s modules=%-3d new=%-3d  (%.0fs)'
                  % (version, len(seen), new, time.time() - t0), flush=True)
        except Exception as e:
            state[version] = {'error': '%s: %s' % (type(e).__name__, str(e)[:80])}
            print('%-9s ERROR %s' % (version, str(e)[:60]), flush=True)
        finally:
            for p in (exe,):
                if os.path.exists(p):
                    os.remove(p)
            shutil.rmtree(xdir, ignore_errors=True)
            json.dump(state, open(state_p, 'w'), indent=1)
            done += 1
    print('done: %d processed this run, %d total known bodies'
          % (done, len(known)), flush=True)
    return 0


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