#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# 画像ファイルと.captionと.tagsファイルを.txtファイルに連結する
# Concatenate image files and .caption and .tags files to .txt files
#
# This script walks through a directory, identifies image files, and checks for the existence of corresponding
# .caption and .tags files. It then concatenates the contents of .caption and .tags files into the .txt files.
#
# Usage:
#     - Place the script in the directory containing the image files.
#     - Run the script to concatenate .caption and .tags files into .txt files.
#     - Use the dry_run flag to preview the changes without writing to the .txt files.
#
# Functions:
#     get_files(path): Walks through the directory and yields image files along with their .caption and .tags files.
#     concat(caption_path, tags_path, txt_path, dry_run=False): Concatenates the contents of .caption and .tags files into the .txt file.

from pathlib import Path
import os
import re
from utils.file_processor import FileProcessor, ProcessorOptions

FILE_EXTS = {".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif", ".jxl"}

def get_files(path):
    path = Path(path)
    # Walk the directory, looking for image files
    for root, dirs, files in os.walk(path):
        root = path / root
        for file in files:
            file = root / file
            if file.suffix not in FILE_EXTS:
                continue
            caption = file.with_suffix(".caption")
            tags = file.with_suffix(".tags")
            txt = file.with_suffix(".txt")
            if not caption.exists():
                print(f"{caption} does not exist")
            if not tags.exists():
                print(f"{tags} does not exist")
            yield file, caption, tags, txt

class CaptionSelectorProcessor(FileProcessor):
    def process_content(self, content: str) -> str:
        captions = re.split(r'----------', content)
        captions = [caption.replace('\n', ' ').strip() for caption in captions if caption.strip()]
        
        best_caption = ""
        for caption in captions:
            if caption and caption[-1] in ".!?":
                if len(caption) > len(best_caption):
                    best_caption = caption
        
        return best_caption

class CaptionConcatenator:
    def __init__(self, dry_run: bool = False):
        self.dry_run = dry_run
        self.caption_processor = CaptionSelectorProcessor(
            ProcessorOptions(recursive=False, dry_run=dry_run, file_extensions={'.caption'})
        )
    
    def concat(self, caption_path: Path, tags_path: Path, txt_path: Path):
        with open(tags_path, "r") as f:
            tags = f.read().strip(", \n")
        
        best_caption = self.caption_processor.process_content(caption_path.read_text())
        if not best_caption:
            self.caption_processor.logger.warning(f"No suitable caption found in {caption_path}")
            return
        
        txt = f"{tags}, {best_caption}"
        
        if self.dry_run:
            print(f"{txt_path}:")
            print(txt)
            print()
        else:
            txt_path.write_text(txt)
            print(f"wrote {txt_path}")

def main():
    dry_run = False
    concatenator = CaptionConcatenator(dry_run=dry_run)
    
    for f in get_files("."):
        concatenator.concat(*f[1:])

if __name__ == "__main__":
    main()