"""
Creates a meta-data file by combining the information from directory structure.
"""
import json
from pathlib import Path

from pydantic import Field

from realfake.config import IMAGE_FORMATS
from realfake.utils import inject_args, Args


class CreateMetadataArgs(Args):
    root_dir: Path
    datasets: str = Field(..., help="Comma-separated list of datasets to include in the meta-data file")
    jsonl_file: Path = Field(..., help="Path to the output JSONL file")


@inject_args
def main(args: CreateMetadataArgs) -> None:
    datasets = args.datasets.split(",")
    records = []
    for dataset in datasets:
        label = "real" if dataset.startswith("real") else "fake"
        dirpath = args.root_dir/dataset
        assert dirpath.exists(), f"dataset dir does not exist: {dirpath}"
        records.extend((parse_imagenet if "imagenet" in dataset else parse_flat)(dirpath, label))
    with open(args.jsonl_file, "w") as f:
        for record in records:
            f.write(json.dumps(record) + "\n")


def parse_imagenet(dirpath: Path, label: str) -> list:
    records = []
    for classdir in dirpath.iterdir():
        assert classdir.is_dir(), f"class directory is not a directory: {classdir}"
        for fn in classdir.iterdir():
            if fn.suffix.lower() in IMAGE_FORMATS:
                records.append({"path": str(fn), "label": label, "class": classdir.name})
            else:
                print("Not an image file:", fn)
    return records


def parse_flat(dirpath: Path, label: str) -> list:
    records = []
    for fn in dirpath.iterdir():
        if fn.suffix.lower() in IMAGE_FORMATS:
            records.append({"path": str(fn), "label": label, "class": None})
        else:
            print("Not an image file:", fn)
    return records



if __name__ == "__main__":
    main()