#!/usr/bin/env python3
"""Make finished tasks pending again by removing their rows from results.jsonl.

The queue server treats any id in results.jsonl as done (except 'ratelimited'), so a task
that failed is never retried automatically — a timeout may have burned credits invisibly
and we don't re-bill by accident. This is the deliberate, explicit way to retry.

    python3 requeue.py --list                  # what's failed, grouped
    python3 requeue.py --dry cal79_            # what WOULD be requeued (prefix match)
    python3 requeue.py cal79_                  # do it (writes a .bak first)
    python3 requeue.py --failed                # every failed task
    python3 requeue.py cal79_disco cal79_tmi   # exact ids
    python3 requeue.py --force cal79_throbbing # re-run a SUCCESS (it drew the wrong thing)

Successes are protected: re-running one spends credits for an asset you already have.
--force is the deliberate override, for when Flow succeeded at the wrong picture.

Stop the queue server first, or it may rewrite the file under you.
"""
import collections
import json
import os
import shutil
import sys

HERE = os.path.expanduser('~/Library/Application Support/flowrinse')
RESULTS = os.path.join(HERE, 'results.jsonl')


def rows():
    with open(RESULTS) as f:
        return [json.loads(l) for l in f if l.strip()]


def main():
    args = [a for a in sys.argv[1:]]
    if not args:
        print(__doc__)
        return 1
    dry = '--dry' in args
    force = '--force' in args
    only_failed = '--failed' in args
    pats = [a for a in args if not a.startswith('--')]
    rs = rows()

    if '--list' in args:
        c = collections.Counter((r['status'], r['id'].rsplit('_', 1)[0]) for r in rs)
        for (st, grp), n in sorted(c.items()):
            print(f'{st:<12} {grp:<24} {n}')
        return 0

    def hit(r):
        if r['status'] == 'success' and not force:
            return False                      # never re-run a success; that's just burning credits
        if only_failed:
            return r['status'] == 'failed'
        return any(r['id'] == p or r['id'].startswith(p) for p in pats)

    doomed = [r for r in rs if hit(r)]
    if not doomed:
        print('nothing matched (successes are never requeued — pass --force if Flow drew the wrong thing)')
        return 1
    if force and any(r['status'] == 'success' for r in doomed):
        print('--force: re-running SUCCESSES. This spends credits for assets you already have.')
    for r in doomed:
        print(f"  requeue  {r['id']:<24} {r['status']:<10} {str(r.get('error',''))[:46]}")
    print(f"\n{len(doomed)} task(s){'  [DRY RUN]' if dry else ''}")
    if dry:
        return 0

    shutil.copy2(RESULTS, RESULTS + '.bak')
    keep = [r for r in rs if r not in doomed]
    with open(RESULTS + '.tmp', 'w') as f:
        for r in keep:
            f.write(json.dumps(r) + '\n')
    os.replace(RESULTS + '.tmp', RESULTS)     # a kill mid-write can't truncate the real file
    print(f'wrote {len(keep)} rows (backup: results.jsonl.bak). Restart the queue server; '
          'they are pending again.')
    return 0


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