"""Reproduce the disclosed inputs; never overwrites original published outputs."""
import argparse
from dataclasses import asdict
import hashlib
import importlib.metadata
import json
from pathlib import Path
import time


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


def validate(bundle):
    bundle = Path(bundle).resolve()
    data = json.loads((bundle/'inputs.json').read_text(encoding='utf-8'))
    seen = set()
    for sample in data['samples']:
        path = (bundle/sample['path']).resolve()
        if not path.is_relative_to(bundle) or (sample['id'],sample['condition']) in seen:
            raise ValueError('Invalid or duplicate sample path')
        seen.add((sample['id'],sample['condition']))
        if digest(path) != sample['sha256']:
            raise ValueError('Audio hash mismatch: '+sample['path'])
    if len(seen) != 49:
        raise ValueError('Expected 49 frozen audio inputs')
    return data


def transcribe_file(model, path, options):
    start = time.perf_counter()
    segments, info = model.transcribe(str(path), **options)
    segments = list(segments)
    return dict(text=' '.join(s.text.strip() for s in segments),
                segments=[asdict(s) for s in segments], duration=info.duration,
                elapsed_seconds=time.perf_counter()-start)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--bundle',type=Path,default=Path(__file__).resolve().parent)
    parser.add_argument('--verify-only',action='store_true')
    parser.add_argument('--model',choices=['tiny.en','base.en'],default='tiny.en')
    parser.add_argument('--output',type=Path,default=Path('my-transcripts.jsonl'))
    args = parser.parse_args()
    data = validate(args.bundle)
    if args.verify_only:
        print(json.dumps({'verified_audio_files':49,'inference_run':False}))
        return
    if args.output.exists():
        raise FileExistsError('Choose a new output path; existing results are never overwritten')
    for name, version in data['versions'].items():
        if importlib.metadata.version(name) != version:
            raise RuntimeError(f'Use the recorded package version: {name}=={version}')
    from faster_whisper import WhisperModel
    from huggingface_hub import snapshot_download
    config = next(m for m in data['models'] if m['name']==args.model)
    folder = Path(snapshot_download(config['repo'],revision=config['revision'],token=False,
                                   allow_patterns=list(config['files'])))
    for name, expected in config['files'].items():
        if digest(folder/name) != expected:
            raise ValueError('Model file mismatch: '+name)
    model = WhisperModel(str(folder),local_files_only=True,**data['runtime'])
    with args.output.open('x',encoding='utf-8') as stream:
        for sample in data['samples']:
            result = dict(id=sample['id'],condition=sample['condition'],model=args.model,
                          sample_sha256=sample['sha256'])
            try:
                result.update(status='ok',**transcribe_file(model,args.bundle/sample['path'],data['options']))
            except Exception as exc:
                result.update(status='error',error=f'{type(exc).__name__}: {exc}')
            stream.write(json.dumps(result,ensure_ascii=True)+'\n')
            stream.flush()
            print(sample['id'],sample['condition'],result['status'])


if __name__ == '__main__':
    main()