"""Run one offline extract call and compare with the packaged source-oracle fixture."""
import argparse
import json
import os
from pathlib import Path
import sys

sys.dont_write_bytecode = True
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['TRANSFORMERS_OFFLINE'] = '1'
os.environ['TOKENIZERS_PARALLELISM'] = 'false'


def main():
    assets = Path(__file__).resolve().parents[1] / 'host_assets'
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--seq', type=int, choices=(128, 256, 512), default=None,
                        help='Validation override; normally the runtime chooses the smallest window')
    parser.add_argument('--table', choices=('fp32', 'fp16'), default='fp32')
    parser.add_argument('--result', type=Path)
    args = parser.parse_args()
    sys.path.insert(0, str(assets))
    import torch
    from runtime import HostRuntime, compare_entities
    from runtime.host_runtime import LABELS
    torch.set_num_threads(4)
    example = json.loads((assets/'example.json').read_text())
    with HostRuntime(assets, table=args.table) as host:
        if args.seq is None:
            entities = host.extract(example['text'], LABELS)
            details = host.last_result
        else:
            details = host.extract_with_details(example['text'], LABELS, seq=args.seq)
            entities = details['entities']
        comparison = compare_entities(entities, example['entities'])
    passed = details['all_finite'] and comparison['span_sets_equal'] and comparison['max_score_diff'] <= 1e-2
    capture = details['captured']
    result = {'status': 'PASS' if passed else 'FAIL', 'text': example['text'], 'entities': entities,
              'comparison': comparison, 'score_atol': 1e-2, 'selected_shape': details['selected_shape'],
              'graph_roles': list(details['backends']), 'backends': details['backends'],
              'embedding_table': args.table, 'encoded_length': capture['encoded_length'],
              'text_word_length': capture['text_word_length'], 'all_finite': details['all_finite'],
              'HF_HUB_OFFLINE': os.environ['HF_HUB_OFFLINE']}
    if args.result:
        args.result.write_text(json.dumps(result, indent=2)+'\n')
    print(json.dumps(result, indent=2))
    return 0 if passed else 2


if __name__ == '__main__':
    raise SystemExit(main())