#!/usr/bin/env python3
# scripts/deck-sgdb-art.py
#
# Batch-download SteamGridDB artwork for RomM-synced shortcuts on the Steam Deck.
# Designed to run ON the Deck (not from Windows) — upload via deck-sgdb-art.ps1.
#
# Usage (on Deck):
#   python3 deck-sgdb-art.py [--types hero,wide,poster] [--clean] [--dry-run]
#
#   --types   Comma-separated list of art types to fetch (default: hero,wide,poster)
#               hero    = background banner  → <appid>_hero.{png,jpg}  (1920x620)
#               wide    = landscape capsule  → <appid>.{png,jpg}        (460x215 / 920x430)
#               poster  = portrait capsule   → <appid>p.{png,jpg}       (600x900)
#   --clean   Remove art for games where no exact SGDB match was found (first-result
#             fallbacks can be misleading for ROM hacks). Safe to run after a bulk pass.
#   --dry-run Print what would happen without downloading or deleting anything.
#
# Gotchas:
#   * Cloudflare blocks Python's default urllib UA — browser UA is required.
#   * appid formula: (raw_appid & 0xffffffff) | 0x80000000
#   * vdf module lives at ~/homebrew/plugins/decky-steamgriddb/py_modules
#   * All RomM shortcuts use 'rom-launcher' in the Exe path.
#   * Steam user for shortcuts/grid: 43872485 (check userdata/ if this ever changes)
#   * SGDB_KEY env var or hardcoded below.

import sys, os, time, json, re
sys.path.insert(0, os.path.expanduser('~/homebrew/plugins/decky-steamgriddb/py_modules'))
from vdf import binary_load
from urllib.request import Request, urlopen
from urllib.parse import quote
from urllib.error import HTTPError, URLError
from pathlib import Path

SGDB_KEY  = os.environ.get('SGDB_KEY', '25431ca935008934cae436d78e9d451c')
STEAM_USER = '43872485'
GRID_DIR   = Path(f'/home/deck/.local/share/Steam/userdata/{STEAM_USER}/config/grid')
VDF_PATH   = Path(f'/home/deck/.local/share/Steam/userdata/{STEAM_USER}/config/shortcuts.vdf')

BROWSER_UA = ('Mozilla/5.0 (X11; Linux x86_64) '
              'AppleWebKit/537.36 (KHTML, like Gecko) '
              'Chrome/120.0.0.0 Safari/537.36')

# --- arg parsing -----------------------------------------------------------
args = sys.argv[1:]
DRY_RUN = '--dry-run' in args
DO_CLEAN = '--clean' in args
types_arg = next((a.split('=', 1)[1] for a in args if a.startswith('--types=')),
                 next((args[args.index('--types') + 1] for i, a in enumerate(args)
                       if a == '--types' and i + 1 < len(args)), None) if '--types' in args else None)
TYPES = set((types_arg or 'hero,wide,poster').split(','))

# --- helpers ---------------------------------------------------------------
def sgdb_get(path):
    req = Request(f'https://www.steamgriddb.com/api/v2{path}',
                  headers={'Authorization': f'Bearer {SGDB_KEY}', 'User-Agent': BROWSER_UA})
    try:
        r = urlopen(req, timeout=15)
        return json.loads(r.read())
    except HTTPError as e:
        return {'success': False, 'error': f'HTTP {e.code}'}
    except URLError as e:
        return {'success': False, 'error': str(e)}

def download(url, out_path):
    req = Request(url, headers={'User-Agent': BROWSER_UA})
    try:
        data = urlopen(req, timeout=30).read()
        if not DRY_RUN:
            with open(out_path, 'wb') as f:
                f.write(data)
        return True
    except Exception as e:
        print(f'    download error: {e}')
        return False

def clean_name(name):
    name = re.sub(r'\.[a-zA-Z0-9]{2,4}$', '', name)   # strip file extension
    name = re.sub(r'\s*[\(\[].+?[\)\]]', '', name)     # strip (USA), [!], (Hack), etc.
    return name.strip()

def existing(appid, suffix, exts=('png', 'jpg')):
    return next((GRID_DIR / f'{appid}{suffix}.{e}' for e in exts
                 if (GRID_DIR / f'{appid}{suffix}.{e}').exists()), None)

def remove_if_exists(appid, suffix):
    for ext in ('png', 'jpg'):
        p = GRID_DIR / f'{appid}{suffix}.{ext}'
        if p.exists():
            if not DRY_RUN:
                p.unlink()
            print(f'  removed: {p.name}')
            return True
    return False

# --- load shortcuts --------------------------------------------------------
d = binary_load(open(VDF_PATH, 'rb'))
romm = [s for s in d['shortcuts'].values()
        if 'rom-launcher' in s.get('Exe', '') or 'romm' in s.get('Exe', '').lower()]
print(f'RomM shortcuts: {len(romm)}  |  types: {", ".join(sorted(TYPES))}  |  dry_run={DRY_RUN}')

# art type config: (flag_name, grid_path_or_None, file_suffix, dimensions_param)
ART_TYPES = {
    'hero':   ('/heroes/game/{id}',                          '_hero', None),
    'wide':   ('/grids/game/{id}?dimensions=460x215,920x430', '',     None),
    'poster': ('/grids/game/{id}?dimensions=600x900,342x482,660x930', 'p', None),
}

done = 0; skipped = 0; not_found = 0; errors = 0

for i, shortcut in enumerate(romm):
    appid = (shortcut['appid'] & 0xffffffff) | 0x80000000
    name  = shortcut.get('appname', '')

    need = {t for t in TYPES if not existing(appid, ART_TYPES[t][1])}

    if DO_CLEAN:
        # In clean mode, remove art for games whose names look like unmatched ROM hacks
        # (file extension still in name = igir didn't find a clean title)
        if re.search(r'\.[a-zA-Z0-9]{2,4}$', name):
            removed = 0
            for t in TYPES:
                if remove_if_exists(appid, ART_TYPES[t][1]):
                    removed += 1
            if removed:
                print(f'[clean] {name} — removed {removed} file(s)')
            skipped += 1
        continue

    if not need:
        skipped += 1
        continue

    search_name = clean_name(name)
    label = f'[{i+1}/{len(romm)}] {name}'
    if search_name != name:
        label += f" -> '{search_name}'"
    print(label)
    sys.stdout.flush()

    if DRY_RUN:
        print(f'  would fetch: {", ".join(sorted(need))}')
        done += 1
        continue

    try:
        search = sgdb_get(f'/search/autocomplete/{quote(search_name)}')
        if not search.get('success') or not search.get('data'):
            print(f'  NOT FOUND ({search.get("error", "no data")})')
            not_found += 1
            time.sleep(0.2)
            continue

        game    = search['data'][0]
        game_id = game['id']
        print(f'  matched: {game["name"]} (#{game_id})')
        sys.stdout.flush()

        for art_type in sorted(need):
            endpoint, suffix, _ = ART_TYPES[art_type]
            resp = sgdb_get(endpoint.format(id=game_id))
            if resp.get('success') and resp.get('data'):
                url = resp['data'][0]['url']
                ext = url.rsplit('.', 1)[-1].lower().split('?')[0]
                out = GRID_DIR / f'{appid}{suffix}.{ext}'
                if download(url, out):
                    print(f'  {art_type}: {out.name}')
                else:
                    print(f'  {art_type}: download failed')
            else:
                print(f'  {art_type}: no results on SGDB')
            sys.stdout.flush()

        done += 1
        time.sleep(0.35)

    except Exception as e:
        print(f'  ERROR: {e}')
        sys.stdout.flush()
        errors += 1
        time.sleep(1)

print(f'\nDone={done} Skipped={skipped} NotFound={not_found} Errors={errors}')