#!/usr/bin/env python3
"""Verify all bytes, or restore 34 archived completed-unit checkouts into fresh paths."""
import argparse,gzip,hashlib,json,os,pathlib,re,shutil,subprocess,tarfile,tempfile
HERE=pathlib.Path(__file__).resolve().parent

def digest(path):
 h=hashlib.sha256()
 with path.open('rb') as f:
  while block:=f.read(1024*1024):h.update(block)
 return h.hexdigest()
def relative(name):
 p=pathlib.PurePosixPath(name)
 if p.is_absolute() or '..' in p.parts or not p.parts or p.parts[0]=='.git':raise ValueError('unsafe source path')
 return p

def main():
 p=argparse.ArgumentParser(description=__doc__);p.add_argument('--archive',type=pathlib.Path,required=True);p.add_argument('--verify-only',action='store_true');p.add_argument('--destination-base',type=pathlib.Path);p.add_argument('--repo',type=pathlib.Path);p.add_argument('--files-only',action='store_true');p.add_argument('--git-bundles',type=pathlib.Path,default=HERE/'git-bundles.tar.gz');args=p.parse_args()
 prep=json.loads((HERE/'PREPARATION.json').read_text());gm=json.loads((HERE/'GIT_RESTORE.json').read_text());assert digest(args.archive)==prep['archive']['sha256'],'archive SHA mismatch'
 raw=gzip.decompress((HERE/'MANIFEST.json.gz').read_bytes());m=json.loads(raw);references={}
 def destination_for(name):
  root=next(v for v in m['roots'] if v['name']==name)
  return args.destination_base/pathlib.Path(root['path']).parent.name/name
 for root in m['roots']:
  assert root['name'] in {v['name'] for v in prep['roots']} and pathlib.Path(root['path']).name==root['name'] and pathlib.Path(root['path']).parent.name in {'p08-02','p08-03-code','p24-03-code'} and re.fullmatch('[0-9a-f]{40}',root['head'])
  for name,row in root['files'].items():relative(name);references.setdefault(row['sha256'],[]).append((root['name'],name,row));assert row['bytes']>=0
  for name in root['directories']:relative(name)
 def archive_pass(write=False):
  seen=set()
  with tarfile.open(args.archive,'r:gz') as tar:
   for member in tar:
    assert member.isfile() and not member.issym() and not member.islnk()
    f=tar.extractfile(member)
    if member.name=='manifest.json':assert f.read()==raw;continue
    assert re.fullmatch('blobs/[0-9a-f]{64}',member.name);h=member.name.split('/')[1];assert h in references and h not in seen
    data=f.read();assert hashlib.sha256(data).hexdigest()==h and all(len(data)==v[2]['bytes'] for v in references[h]);seen.add(h)
    if write:
     for root,name,row in references[h]:
      target=destination_for(root)/name;target.parent.mkdir(parents=True,exist_ok=True);target.write_bytes(data);target.chmod(row['mode']);os.utime(target,ns=(row['mtime_ns'],row['mtime_ns']))
  assert seen==set(references)
 archive_pass()
 assert sum(len(r['files']) for r in m['roots'])==m['files']
 if args.verify_only:print(json.dumps(dict(verified=True,files=m['files'],logical_bytes=m['logical_bytes'],unique_blobs=len(references))));return
 assert args.destination_base and (args.files_only or args.repo),'choose explicit destination and repo or files-only'
 args.destination_base=args.destination_base.resolve();args.destination_base.mkdir(parents=True,exist_ok=True)
 for root in m['roots']:assert not (destination_for(root['name'])).exists(),'fresh destination required'
 if not args.files_only:
  assert digest(args.git_bundles)==gm['archive']['sha256'],'Git bundles SHA mismatch'
  with tempfile.TemporaryDirectory(prefix='p04-restore-git-') as tmp:
   with tarfile.open(args.git_bundles,'r:gz') as tar:
    for member in tar:
     assert member.isfile() and re.fullmatch('[0-9a-f]{40}.bundle',member.name)
     target=pathlib.Path(tmp)/member.name;target.write_bytes(tar.extractfile(member).read());record=next(r for r in gm['commits'] if r.get('file')==member.name);assert digest(target)==record['sha256']
     subprocess.run(['git','-C',str(args.repo),'bundle','verify',str(target)],check=True,capture_output=True)
     subprocess.run(['git','-C',str(args.repo),'bundle','unbundle',str(target)],check=True,capture_output=True)
  for root in m['roots']:subprocess.run(['git','-C',str(args.repo),'worktree','add','--detach',str(destination_for(root['name'])),root['head']],check=True)
 else:
  for root in m['roots']:(destination_for(root['name'])).mkdir(parents=True)
 archive_pass(write=True)
 for root in m['roots']:
  destination=destination_for(root['name'])
  for name,record in sorted(root['directories'].items(),key=lambda q:len(pathlib.PurePosixPath(q[0]).parts),reverse=True):
   target=destination/name;target.mkdir(parents=True,exist_ok=True);target.chmod(record['mode']);os.utime(target,ns=(record['mtime_ns'],record['mtime_ns']))
  actual={f.relative_to(destination).as_posix() for f in destination.rglob('*') if f.is_file() and f.relative_to(destination).as_posix()!='.git'};assert actual==set(root['files'])
  for name,row in root['files'].items():
   f=destination/name;assert digest(f)==row['sha256'] and f.stat().st_size==row['bytes'] and f.stat().st_mtime_ns==row['mtime_ns']
 (args.destination_base/'RESTORED_COMMITS.json').write_text(json.dumps([{k:r[k] for k in ['name','head']} for r in m['roots']],indent=2)+'\n')
 print(json.dumps(dict(restored=True,files=m['files'],destination=str(args.destination_base))))
if __name__=='__main__':main()