"""Identify evaluation examples and protect saved prediction runs."""
import hashlib
from importlib import metadata
import json
import os
from pathlib import Path


def sample_identifier(record):
    value = record.get('example_id', record.get('pair_id'))
    if value is None:
        raise ValueError('Each annotation must have an example_id or pair_id.')
    value = str(value)
    if not value or value in ('.', '..') or Path(value).name != value or '\\' in value:
        raise ValueError('Sample IDs must be nonempty and contain no path separators.')
    return value


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


def checkpoint_digest(path):
    path = Path(path)
    if path.is_file():
        return file_digest(path)
    suffixes = {'.safetensors', '.bin', '.pth', '.pt', '.json', '.model',
                '.txt', '.py', '.vocab', '.merges'}
    files = sorted(file for file in path.rglob('*')
                   if file.is_file() and file.suffix in suffixes
                   and not any(part.startswith('.') or part == 'code'
                               for part in file.relative_to(path).parts))
    if not files:
        raise ValueError(f'No model or tokenizer assets found in {path}.')
    digest = hashlib.sha256()
    for file in files:
        digest.update(file.relative_to(path).as_posix().encode() + b'\0')
        digest.update(bytes.fromhex(file_digest(file)))
    return digest.hexdigest()


def index_prediction_files(directory):
    files = {}
    for path in sorted(Path(directory).rglob('*.pth')):
        if path.stem in files:
            raise ValueError(
                f'Duplicate predictions for {path.stem}: {files[path.stem]} and {path}. '
                'Keep one prediction per example before scoring.'
            )
        files[path.stem] = path
    return files


def generation_identity(args, checkpoint):
    asset_digests = {}

    def digest_asset(value):
        path = Path(value).resolve()
        if path not in asset_digests:
            asset_digests[path] = checkpoint_digest(path)
        return asset_digests[path]

    settings = ('precision', 'image_size', 'vit_image_size', 'model_max_length',
                'max_num_images', 'num_query_tokens', 'lora_r', 'lora_alpha',
                'lora_dropout', 'lora_target_modules', 'conv_type',
                'use_mm_start_end', 'max_new_tokens', 'm4seg_data',
                'n_chunks', 'chunk_idx')
    identity = {
        'schema_version': 2,
        'annotation_sha256': file_digest(args.m4seg_annpath),
        'checkpoint_sha256': digest_asset(checkpoint),
        'initialization': 'merged' if args.merged_ckpt_path else 'backbone',
        'dataset_root': str(Path(args.dataset_dir).resolve()),
        'settings': {name: getattr(args, name) for name in settings},
        'preprocessing': {
            'truncate_from_left': os.getenv('PRIMA_TRUNCATE_FROM_LEFT', '0').strip().lower()
            in ('1', 'true', 'yes', 'y', 'on'),
        },
    }
    if args.merged_ckpt_path:
        if args.ckpt_path and not args.ckpt_path.startswith('/path/to/'):
            identity['overlay_sha256'] = digest_asset(args.ckpt_path)
    else:
        assets = ('version', 'q_former_model', 'pretrained_q_former_path',
                  'qformer_bert_path', 'eva_weights', 'pretrain_mm_mlp_adapter')
        identity['initialization_assets'] = {
            name: digest_asset(getattr(args, name)) for name in assets
            if getattr(args, name, '') and not str(getattr(args, name)).startswith('/path/to/')
        }
    root = Path(__file__).resolve().parents[1]
    digest = hashlib.sha256()
    for directory in ('model', 'dataset', 'tools', 'eval'):
        for path in sorted((root / directory).rglob('*.py')):
            digest.update(path.relative_to(root).as_posix().encode() + b'\0')
            digest.update(bytes.fromhex(file_digest(path)))
    identity['source_sha256'] = digest.hexdigest()
    dependencies = ('torch', 'torchvision', 'transformers', 'tokenizers', 'peft',
                    'timm', 'numpy', 'Pillow', 'opencv-python', 'opencv-python-headless',
                    'sentencepiece', 'einops')
    identity['dependencies'] = {}
    for name in dependencies:
        try:
            identity['dependencies'][name] = metadata.version(name)
        except metadata.PackageNotFoundError:
            identity['dependencies'][name] = None
    return identity


def prepare_prediction_directory(output, identity, *, resume):
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    provenance = output / 'generation.json'
    existing = list(output.rglob('*.pth'))
    if any(path.parent != output for path in existing):
        raise ValueError('Use a separate output directory for each generation run or shard.')
    if existing and not resume:
        raise ValueError('Use an empty output directory or explicitly enable --resume.')
    if provenance.exists():
        if not resume:
            raise ValueError('This output directory already has generation provenance. '
                             'Use --resume or a new directory.')
        try:
            previous = json.loads(provenance.read_text())
        except (OSError, ValueError) as error:
            raise ValueError('Cannot verify existing generation provenance. Use a new output directory.') from error
        if previous != identity:
            raise ValueError('Resume configuration does not match existing predictions. '
                             'Use a new output directory.')
        return
    if existing:
        raise ValueError('Existing predictions have no generation provenance. '
                         'Use a new output directory; their configuration cannot be verified.')
    try:
        with provenance.open('x') as handle:
            json.dump(identity, handle, indent=2)
            handle.write('\n')
    except FileExistsError as error:
        raise ValueError('Another run has already claimed this output directory.') from error