"""Create reproducible real-portrait splits, excluding benchmark duplicates."""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
import cv2


def pixel_digest(path):
    image = cv2.imread(str(path))
    if image is None:
        raise ValueError(f'Unreadable dataset image: {path}')
    return hashlib.sha256(str(image.shape).encode() + image.tobytes()).hexdigest()


def prepare_split(root: Path, holdout: Path, output: Path, folders=('faces', 'pinterest'), seed=42, exclude_unreadable=False):
    root = root.resolve()
    excluded = {s.strip() for s in holdout.read_text(encoding='utf-8').splitlines() if s.strip()}
    heldout_hashes = set()
    for name in sorted(excluded):
        path = (root / name).resolve()
        if not path.is_relative_to(root):
            raise ValueError('Holdout path escapes dataset root')
        heldout_hashes.add(pixel_digest(path))
    seen, rows, rejected = set(), [], []
    duplicates = 0
    for folder in folders:
        for path in sorted((root / folder).rglob('*')):
            if not path.is_file() or path.suffix.lower() not in {'.png', '.jpg', '.jpeg', '.webp'}:
                continue
            try:
                digest = pixel_digest(path)
            except ValueError:
                if not exclude_unreadable:
                    raise
                rejected.append({'path': path.relative_to(root).as_posix(), 'reason': 'unreadable'})
                continue
            if digest in heldout_hashes or digest in seen:
                duplicates += 1
                continue
            seen.add(digest)
            rows.append({'path': path.relative_to(root).as_posix(), 'pixels_sha256': digest})
    if len(rows) < 100:
        raise ValueError(f'Only {len(rows)} unique non-holdout portraits; refusing a toy split')
    rows.sort(key=lambda row: hashlib.sha256(f"{seed}:{row['pixels_sha256']}".encode()).hexdigest())
    validation_count = max(1, round(len(rows) * .05))
    validation, train = rows[:validation_count], rows[validation_count:]
    if output.exists() and any(output.iterdir()):
        raise FileExistsError('Choose a new split directory to preserve existing splits')
    output.mkdir(parents=True, exist_ok=True)
    for name, items in [('train', train), ('validation', validation)]:
        (output / f'{name}.txt').write_text(''.join(row['path'] + '\n' for row in items), encoding='utf-8')
    report = {'schema_version': 1, 'seed': seed, 'domains': list(folders),
              'holdout_count': len(excluded), 'excluded_or_duplicate_count': duplicates,
              'train': train, 'validation': validation, 'rejected': rejected,
              'review_status': 'pending_identity_and_near_duplicate_review'}
    benchmark_manifest = holdout.parent / 'manifest.csv'
    if benchmark_manifest.is_file():
        report['benchmark_manifest_sha256'] = hashlib.sha256(benchmark_manifest.read_bytes()).hexdigest()
    (output / 'holdout_paths.txt').write_text(holdout.read_text(encoding='utf-8'), encoding='utf-8')
    (output / 'split.json').write_text(json.dumps(report, indent=2), encoding='utf-8')
    print(f'Split: {len(train)} train, {len(validation)} validation; {len(excluded)} holdout sources excluded')
    return report


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--root', type=Path, default=Path('models/CodeFormer/datasets/ffhq/ffhq_512'))
    parser.add_argument('--holdout', type=Path, default=Path('benchmarks/holdout_paths.txt'))
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--exclude-unreadable', action='store_true', help='Record unreadable images as excluded; never delete them')
    args = parser.parse_args()
    prepare_split(args.root, args.holdout, args.output, exclude_unreadable=args.exclude_unreadable)