#!/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

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
    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':
            return False                      # never re-run a success; that's just burning credits
        if only_failed and r['status'] == 'failed':
            return True
        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)')
        return 1
    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())