"""
Check that files that are references from JSONL file are valid.
"""
import json
from pathlib import Path

import numpy as np
import PIL.Image
from joblib import delayed, Parallel

from realfake.utils import inject_args, Args, read_jsonl


class CheckFilesArgs(Args):
    jsonl_file: Path


@inject_args
def main(args: CheckFilesArgs) -> None:
    records = read_jsonl(args.jsonl_file)
    results = Parallel(n_jobs=-1, verbose=100)(delayed(check_file)(record) for record in records)
    failed = [result for result in results if result["error"] is not None]
    if not failed:
        print("All files are valid")
    else:
        saved_file = args.jsonl_file.with_suffix(".failed.jsonl")
        print(f"{len(failed)} files are invalid, saved errors to {saved_file}")
        with open(saved_file, "w") as f:
            for record in failed:
                f.write(json.dumps(record) + "\n")


def check_file(record: dict) -> dict:
    path = Path(record["path"])
    error = None
    if not path.exists():
        error = "File does not exist"
    elif not path.is_file():
        error = "Path is not a file"
    elif path.suffix.lower() not in (".jpg", ".jpeg", ".png"):
        error = "File is not an image file"
    else:
        try:
            np.asarray(PIL.Image.open(path))
        except Exception as e:
            error = f"Image cannot be opened: {e}"
    return dict(record, error=error)
            

if __name__ == '__main__':
    main()