#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# ディレクトリを走査し、画像ファイルを特定し、対応する.caption、.wd、.tagsファイルの存在を確認します。
# その後、これらのファイルの内容を.txtファイルに連結します。
# This script walks through a directory, identifies image files, and checks for the existence of corresponding
# .caption, .wd, and .tags files. It then concatenates the contents of these files into a single line in the .txt file.
#
# Usage:
#     - Place the script in the directory containing the image files.
#     - Run the script to concatenate .caption, .wd, and .tags files into .txt files.
#     - Use the dry_run flag to preview the changes without writing to the .txt files.
#     - Files missing any of .caption, .wd, or .tags will be skipped.
#
# Functions:
#     get_files(path): Walks through the directory and yields image files along with their .caption, .wd, and .tags files.
#     concat_files_single_line(caption_path, wd_path, tags_path, txt_path, dry_run=False): Concatenates the contents into a single line in the .txt file.

from pathlib import Path
import os
import argparse

# Supported image file extensions
FILE_EXTS = {".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif", ".jxl", ".webp", ".avif"}


def get_files(path):
    """
    Walks through the directory and yields image files along with their corresponding .caption, .wd, and .tags files.

    :param path: Path to the directory to search.
    :yield: Tuple containing image file, caption file, wd file, tags file, and txt file paths.
    """
    path = Path(path)
    for root, dirs, files in os.walk(path):
        root = Path(root)
        for file in files:
            file_path = root / file
            if file_path.suffix.lower() not in FILE_EXTS:
                continue

            caption = file_path.with_suffix(".caption")
            wd = file_path.with_suffix(".wd")
            tags = file_path.with_suffix(".tags")
            txt = file_path.with_suffix(".txt")

            missing = []
            if not caption.exists():
                missing.append(caption.name)
            if not wd.exists():
                missing.append(wd.name)
            if not tags.exists():
                missing.append(tags.name)

            if missing:
                print(f"Skipping {file_path.name}: Missing files: {', '.join(missing)}")
                continue

            yield file_path, caption, wd, tags, txt


def concat_files_single_line(caption_path, wd_path, tags_path, txt_path, dry_run=False):
    """
    Concatenates the contents of .caption, .wd, and .tags files into a single line in a .txt file.

    :param caption_path: Path to the .caption file.
    :param wd_path: Path to the .wd file.
    :param tags_path: Path to the .tags file.
    :param txt_path: Path to the output .txt file.
    :param dry_run: If True, prints the concatenated content without writing to the file.
    """
    try:
        with open(caption_path, "r", encoding="utf-8") as f:
            caption = f.read().strip()

        with open(wd_path, "r", encoding="utf-8") as f:
            wd = f.read().strip(", \n")

        with open(tags_path, "r", encoding="utf-8") as f:
            tags = f.read().strip(", \n")

        wd_tags = [tag.strip() for tag in wd.split(',')]
        tags_tags = [tag.strip() for tag in tags.split(',')]
        unique_tags = []
        for tag in wd_tags + tags_tags:
            if tag and tag not in unique_tags:
                unique_tags.append(tag)
        concatenated = f"{', '.join(unique_tags)}, {caption}"

        if dry_run:
            print(f"{txt_path}:")
            print(concatenated)
            print()
        else:
            with open(txt_path, 'w', encoding="utf-8") as f:
                f.write(concatenated)
            print(f"Wrote {txt_path}")

    except Exception as e:
        print(f"Error processing {txt_path}: {e}")


def main(directory, dry_run=False):
    """
    Main function to process all relevant files in the specified directory.

    :param directory: Directory to process.
    :param dry_run: If True, perform a dry run without writing files.
    """
    for file, caption, wd, tags, txt in get_files(directory):
        concat_files_single_line(caption, wd, tags, txt, dry_run=dry_run)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Concatenate .caption, .wd, and .tags files into .txt files as a single line.")
    parser.add_argument(
        "-d",
        "--directory",
        type=str,
        default=".",
        help="Directory containing the image files. Defaults to the current directory.",
    )
    parser.add_argument(
        "--dry_run",
        action="store_true",
        help="If set, previews the concatenated content without writing to .txt files.",
    )
    args = parser.parse_args()

    main(args.directory, dry_run=args.dry_run)