"""Separate, truthful cache schema for authored audio-training candidates."""
from common import *
from audio_data_candidates import candidates,validate

CACHE=EXP/'audio_candidate_cache'


def candidate_manifest():
    manifest=read(EXP/'audio_data_candidates.json')
    rows=candidates()
    if manifest['candidates']!=rows or manifest['checks']!=validate(rows) or manifest['source_sha256']!=sha(EXP/'audio_data_candidates.py'):
        raise RuntimeError('Authored candidate source/manifest mismatch')
    return manifest


def qualified_teacher():
    root=EXP/'audio_data_qualification'
    receipt=read(root/'SUCCESS.json')
    if receipt['status']!='success' or receipt['samples']!=8 or receipt['speech_samples']!=4:
        raise RuntimeError('Teacher-data preflight is incomplete')
    for filename,key in (('MANIFEST.json','manifest_sha256'),('AUDIO_DIAGNOSTICS.json','diagnostics_sha256'),
                         ('asr/TRAINING_ASR.json','asr_sha256')):
        if sha(root/filename)!=receipt[key]:
            raise RuntimeError('Teacher-data preflight report changed')
    manifest=read(root/'MANIFEST.json')
    from qualify_audio_data import plan as qualification_plan
    if manifest['plan']!=qualification_plan():
        raise RuntimeError('Teacher preflight used different data or Base revision')
    selected={row['id']:row for row in manifest['plan']['samples']}
    if len(manifest['samples'])!=8 or {row['id'] for row in manifest['samples']}!=set(selected):
        raise RuntimeError('Teacher-data media coverage mismatch')
    for row in manifest['samples']:
        if any(row[key]!=selected[row['id']][key] for key in ('prompt','seed')):
            raise RuntimeError('Teacher-data media request mismatch')
        if sha(row['video'])!=row['video_sha256'] or sha(row['native_wav'])!=row['native_wav_sha256']:
            raise RuntimeError('Teacher-data media changed')
    asr=read(root/'asr/TRAINING_ASR.json')
    speech={key:row for key,row in selected.items() if row['phrase']}
    if len(asr['results'])!=4 or {r['case'] for r in asr['results']}!=set(speech) or any(
            r['edit_distance']!=0 or r['text_prompt_supplied_to_asr'] or r['reference']!=speech[r['case']]['phrase']
            or r['language']!=speech[r['case']]['language'] for r in asr['results']):
        raise RuntimeError('Sampled teacher transcripts require review before encoding this trial')
    if not receipt['sampled_speech_reference_gate_passed'] or receipt['speech_exact_matches']!=4:
        raise RuntimeError('Sampled teacher speech-reference gate has not passed')
    return sha(root/'SUCCESS.json')


def contract():
    candidate_manifest()
    source_contract()
    return {'artifact_type':'h3_authored_audio_candidate_embeddings_v1','schema_version':1,
        'model_revision':REVISION,'workflow':'positive_text_encoding','IR':False,
        'candidate_manifest_sha256':sha(EXP/'audio_data_candidates.json'),
        'qualified_teacher_receipt_sha256':qualified_teacher(),
        'qualification_scope':'Sampled transcript gate only, not perceptual or whole-pool quality.',
        'sources':{name:sha(EXP/name) for name in ('audio_candidate_cache.py','encode_audio_candidates.py','audio_data_candidates.py')},
        'official_encoder_sha256':sha(ASSETS/'external/diffusers-minimax-h3/src/diffusers/modular_pipelines/minimax_h3/encoders.py'),
        'text_dim':read(BASE/'transformer/config.json')['text_dim']}


def validate_condition(condition,row):
    imports()
    from minimax_h3.encode_longlive_prompt_cache_v4 import _validate_encoded_condition
    expected={'id':row['id'],'prompt':row['prompt'],
              'prompt_sha256':hashlib.sha256(row['prompt'].encode()).hexdigest()}
    if any(condition.get(key)!=value for key,value in expected.items()):
        raise RuntimeError('Encoded candidate identity mismatch')
    _validate_encoded_condition(condition)
    if condition['prompt_embeds'].shape[-1]!=read(BASE/'transformer/config.json')['text_dim']:
        raise RuntimeError('Candidate embedding width differs from original Base')


def load_condition(row,expected_contract):
    import torch
    root=CACHE/row['id']
    receipt=read(root/'SUCCESS.json')
    path=root/'condition.pt'
    if receipt['status']!='success' or receipt['contract']!=expected_contract or receipt['id']!=row['id']:
        raise RuntimeError('Candidate encoding receipt mismatch')
    if path.stat().st_size!=receipt['bytes'] or sha(path)!=receipt['sha256']:
        raise RuntimeError('Candidate tensor payload changed')
    payload=torch.load(path,map_location='cpu',weights_only=True)
    if payload['contract']!=expected_contract:
        raise RuntimeError('Candidate tensor contract mismatch')
    condition=payload['condition']
    validate_condition(condition,row)
    return condition


def verify_collection():
    expected=contract()
    manifest=read(CACHE/'COLLECTION.json')
    rows=candidate_manifest()['candidates']
    if manifest['contract']!=expected or [r['id'] for r in manifest['conditions']]!=[r['id'] for r in rows]:
        raise RuntimeError('Audio candidate cache coverage/contract mismatch')
    result={}
    for row,record in zip(rows,manifest['conditions'],strict=True):
        if record['receipt_sha256']!=sha(CACHE/row['id']/'SUCCESS.json'):
            raise RuntimeError('Audio candidate row receipt changed')
        result[row['id']]=load_condition(row,expected)
    return result,{'collection_sha256':sha(CACHE/'COLLECTION.json'),'contract':expected}


def load_pool():
    receipt=read(CACHE/'SUCCESS.json')
    if receipt['status']!='success' or receipt['conditions']!=256 or receipt['collection_sha256']!=sha(CACHE/'COLLECTION.json'):
        raise RuntimeError('Audio candidate cache is incomplete or changed')
    return verify_collection()