#!/usr/bin/env python3
"""Stream a prepared board release to a new local archive using parallel gzip."""
import argparse
import hashlib
import json
from pathlib import Path
import shlex
import subprocess


def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--remote',required=True)
    p.add_argument('--directory',type=Path,required=True)
    p.add_argument('--output',type=Path,required=True)
    a=p.parse_args()
    partial=a.output.with_suffix(a.output.suffix+'.building')
    checksum=a.output.with_name(a.output.name+'.sha256')
    if any(x.exists() for x in (partial,a.output,checksum)):raise FileExistsError('archive output already exists')
    a.output.parent.mkdir(parents=True,exist_ok=True)
    command='tar -C '+shlex.quote(str(a.directory.parent))+' -cf - -- '+shlex.quote(a.directory.name)
    with partial.open('xb') as f:
        remote=subprocess.Popen(['ssh','-o','BatchMode=yes',a.remote,command],stdout=subprocess.PIPE)
        compressor=None
        try:
            compressor=subprocess.Popen(['pigz','-6','-p','8','-c'],stdin=remote.stdout,stdout=f)
            remote.stdout.close()
            if compressor.wait()!=0:raise RuntimeError('compression failed')
            if remote.wait()!=0:raise RuntimeError('remote tar stream failed')
        finally:
            for child in (compressor,remote):
                if child is not None and child.poll() is None:child.terminate();child.wait()
    partial.rename(a.output)
    with a.output.open('rb') as f:sha=hashlib.file_digest(f,'sha256').hexdigest()
    checksum.write_text(sha+'  '+a.output.name+'\n')
    print(json.dumps(dict(archive=str(a.output),bytes=a.output.stat().st_size,sha256=sha)),flush=True)


if __name__=='__main__':main()