"""Resumable KNP detector collection; fixed model/data/build signature, bounded processes."""
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import fcntl
import hashlib
import json
import os
from pathlib import Path
import subprocess
import time


def digest(path):
    h = hashlib.sha256()
    with path.open('rb') as f:
        for block in iter(lambda: f.read(1024 * 1024), b''):
            h.update(block)
    return h.hexdigest()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('model', type=Path)
    ap.add_argument('out', type=Path)
    ap.add_argument('chunk', nargs='?', type=int, default=1000)
    ap.add_argument('--data', type=Path, default=Path.home() / 'run/cifar/vehicles')
    ap.add_argument('--build', type=Path, default=Path(__file__).resolve().parents[1] / 'build')
    ap.add_argument('--backend', choices=['cpu','cuda'], default='cpu')
    ap.add_argument('--jobs', type=int)
    ap.add_argument('--limit-per-split', type=int, help='Diagnostic subset only; use a separate output directory')
    a = ap.parse_args()
    if a.jobs is None: a.jobs = int(os.environ.get("JOBS", "2" if a.backend == "cuda" else "12"))
    if a.chunk < 1 or not 1 <= a.jobs <= 64 or (a.limit_per_split is not None and a.limit_per_split < 1):
        ap.error('Invalid chunk/jobs/limit')
    a.model, a.out, a.data, a.build = (p.resolve() for p in (a.model, a.out, a.data, a.build))
    a.out.mkdir(parents=True, exist_ok=True)
    lock = (a.out / '.lock').open('w')
    fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    executable = a.build / 'bin/altai_infer'
    backend = a.build / 'bin' / ('libknp-altai-cuda-backend.so' if a.backend == 'cuda' else 'libknp-cpu-single-threaded-backend.so')
    k, stride, threshold, classes, _ = map(int, (a.model / 'meta.txt').read_text().split())
    jobs, sizes = [], {}
    for split in ('train', 'train_flip', 'val'):
        image_file = a.data / f'{split}_images.bin'
        size = image_file.stat().st_size
        if size % 3072:
            raise ValueError(f'Invalid image file: {image_file}')
        n = size // 3072
        if a.limit_per_split:
            n = min(n, a.limit_per_split)
        sizes[split] = n
        for first in range(0, n, a.chunk):
            jobs.append((split, first, min(a.chunk, n - first)))
    signature = {'detector': [k, stride, threshold, classes], 'data_sizes': sizes, 'chunk': a.chunk,
                 'files': {str(p): digest(p) for p in [executable, backend, a.model/'det_w.bin', a.model/'det_b.bin',
                                                        *(a.data/f'{s}_images.bin' for s in sizes)]}}
    settings = a.out / 'collection-settings.json'
    if settings.exists() and json.loads(settings.read_text()) != signature:
        raise ValueError('Different model/data/build/chunk: use a new output directory')
    settings.write_text(json.dumps(signature, indent=2) + '\n')
    def valid(job):
        split, first, count = job
        prefix = a.out / f'{split}_{first}'
        output = Path(str(prefix) + '_counts.bin')
        return Path(str(prefix) + '.done').exists() and output.exists() and output.stat().st_size == count * 4 * k * 4
    finished = sum(n for s, first, n in jobs if valid((s, first, n)))
    total = sum(sizes.values())
    started = time.time()
    def progress(failed=False, error=None):
        state = {'current': finished, 'total': total, 'done': finished == total and not failed,
                 'failed': failed, 'error': error, 'started_unix': started, 'updated_unix': time.time()}
        temporary = a.out/'progress.tmp'
        temporary.write_text(json.dumps(state) + '\n')
        temporary.replace(a.out/'progress.json')
    def one(job):
        split, first, count = job
        prefix = a.out / f'{split}_{first}'
        env = dict(os.environ, LD_LIBRARY_PATH=str(a.build/'bin') + ':' + os.environ.get('LD_LIBRARY_PATH',''))
        with Path(str(prefix) + '.log').open('w') as log:
            result = subprocess.run([str(executable), str(a.data/f'{split}_images.bin'), str(first), str(count),
                                     str(a.model), str(backend), str(prefix)], stdout=log, stderr=subprocess.STDOUT, env=env)
        output = Path(str(prefix) + '_counts.bin')
        if result.returncode or not output.exists() or output.stat().st_size != count * 4 * k * 4:
            raise RuntimeError(f'{split}_{first}: failed or truncated output; inspect its log')
        Path(str(prefix) + '.done').write_text(json.dumps({'count':count, 'completed_unix':time.time()}) + '\n')
        return count
    progress()
    errors = []
    with ThreadPoolExecutor(max_workers=a.jobs) as pool:
        futures = {pool.submit(one, job): job for job in jobs if not valid(job)}
        for future in as_completed(futures):
            try:
                finished += future.result()
            except Exception as error:
                errors.append(str(error))
            progress(bool(errors), '; '.join(errors) if errors else None)
            print(f'{finished}/{total}' + (' FAILED '+errors[-1] if errors else ''), flush=True)
    if errors:
        raise RuntimeError('; '.join(errors))


if __name__ == '__main__':
    main()