"""Build a full accession snapshot in resumable parallel partitions.

Run list_annotation_objects.py first. This command indexes the captured inventory,
merges all partitions, and leaves publication to the operator.
"""
import argparse
import json
from pathlib import Path
import sqlite3
import subprocess
import sys

ROOT=Path(__file__).resolve().parent


def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--inventory',type=Path,required=True)
    p.add_argument('--cache-dir',type=Path,default=ROOT/'.cache/full-index')
    p.add_argument('--output',type=Path,default=ROOT/'data/catalog.sqlite')
    p.add_argument('--parts',type=int,default=8)
    p.add_argument('--workers-per-part',type=int,default=32)
    args=p.parse_args()
    if not 1<=args.parts<=32 or not 1<=args.workers_per_part<=128: p.error('Invalid concurrency')
    work=args.cache_dir.resolve(); work.mkdir(parents=True,exist_ok=True)
    config=work/'partition-config.json'
    if config.exists() and json.loads(config.read_text())['parts']!=args.parts:
        p.error('Use the same partition count when resuming, or a new cache directory.')
    config.write_text(json.dumps({'parts':args.parts}))
    items=[json.loads(s) for s in args.inventory.open()]
    expected={x['path']:x['hash'] for x in items}
    seen=set(); assigned=[[] for _ in range(args.parts)]
    for i in range(args.parts):
        done=work/f'part-{i}.sqlite'; checkpoint=work/f'part-{i}.building.sqlite'
        if done.exists() and not checkpoint.exists(): done.rename(checkpoint)
        if checkpoint.exists():
            with sqlite3.connect(checkpoint) as c:
                retained={r[0] for r in c.execute('SELECT path,hash FROM files') if expected.get(r[0])==r[1]}
            if retained & seen: raise ValueError('Duplicate checkpoint objects')
            assigned[i]=[x for x in items if x['path'] in retained]
            seen.update(retained)
    for n,item in enumerate(x for x in items if x['path'] not in seen):
        assigned[n%args.parts].append(item)
    processes=[]
    for i in range(args.parts):
        inventory=work/f'inventory-{i}.jsonl'
        inventory.write_text(''.join(json.dumps(x)+'\n' for x in assigned[i]))
        log=(work/f'part-{i}.log').open('w')
        command=[sys.executable,str(ROOT/'build_full_index.py'),'--inventory',str(inventory),
                 '--output',str(work/f'part-{i}.sqlite'),'--workers',str(args.workers_per_part)]
        processes.append((i,subprocess.Popen(command,stdout=log,stderr=subprocess.STDOUT),log))
    failed=[]
    for i,process,log in processes:
        status=process.wait(); log.close()
        if status: failed.append(i)
    if failed: raise RuntimeError(f'Partitions {failed} failed; inspect logs and rerun to resume.')
    subprocess.run([sys.executable,str(ROOT/'merge_full_index.py'),'--inventory',str(args.inventory.resolve()),
                    '--parts',*[str(work/f'part-{i}.sqlite') for i in range(args.parts)],'--output',str(args.output.resolve())],check=True)


if __name__=='__main__': main()