#!/usr/bin/env python3
"""Predict real/fake for one image or a directory, without torchrun."""
import argparse
import json
import os
from pathlib import Path
import sys

from ltf_final import LTFDetector
from ltf_final.detector import ROOT, BRANCHES
from utils.tile_dataset_srm import IMG_EXTENSIONS


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input', type=Path, required=True, help='Image file or directory, scanned recursively')
    parser.add_argument('--output', type=Path, help='JSONL predictions; stdout is used when omitted')
    parser.add_argument('--config', type=Path, default=ROOT / 'configs/eval.yaml')
    parser.add_argument('--weights', type=Path, help='Override the bundled epoch_20.pth')
    parser.add_argument('--device', default='auto', help='auto, cpu, cuda:0, etc.')
    parser.add_argument('--branch', choices=BRANCHES, default='fusion')
    parser.add_argument('--threshold', type=float, default=.5)
    parser.add_argument('--chunk-size', type=int, default=256, help='CLIP views per forward chunk; tile coverage is preserved')
    parser.add_argument('--threads', type=int, default=4, help='PyTorch CPU thread count')
    args = parser.parse_args()
    if args.threads < 1:
        parser.error('--threads must be positive')
    if not 0 <= args.threshold <= 1:
        parser.error('--threshold must lie in [0,1]')
    source = args.input.expanduser().resolve()
    if source.is_file():
        paths = [source]
    elif source.is_dir():
        paths = sorted(p for p in source.rglob('*') if p.is_file() and p.suffix.lower() in IMG_EXTENSIONS)
    else:
        parser.error(f'Input does not exist: {source}')
    if not paths:
        parser.error('No images found')
    import torch
    torch.set_num_threads(args.threads)
    detector = LTFDetector(args.config, args.weights.resolve() if args.weights else None,
                           args.device, args.chunk_size)
    stream, temporary = sys.stdout, None
    if args.output:
        output = args.output.expanduser().resolve()
        output.parent.mkdir(parents=True, exist_ok=True)
        temporary = output.with_name(output.name + f'.{os.getpid()}.tmp')
        stream = temporary.open('w', encoding='utf-8')
    try:
        for path in paths:
            result = detector.predict(path, args.branch, args.threshold)
            stream.write(json.dumps(result, ensure_ascii=False, allow_nan=False) + '\n')
        stream.flush()
        if temporary:
            stream.close()
            os.replace(temporary, output)
    except BaseException:
        if temporary:
            stream.close()
            temporary.unlink(missing_ok=True)
        raise
    print(f'Predicted {len(paths)} image(s); branch={args.branch}; threshold={args.threshold}', file=sys.stderr)


if __name__ == '__main__':
    main()