import os
import cv2
import pandas as pd
import re

class ImageProcessor:
    def __init__(self):
        self.output_dir = "images-fix"
        os.makedirs(self.output_dir, exist_ok=True)
        
    def get_new_filename(self, original_filename, new_pos_number):
        """ファイル名のpos番号を変更する"""
        pattern = r'(.*pos)(\d+)(.*)'
        return re.sub(pattern, f'\\g<1>{new_pos_number}\\g<3>', original_filename)
    
    def check_duplicate_filename(self, filename):
        """出力ディレクトリで重複をチェック"""
        return os.path.exists(os.path.join(self.output_dir, filename))
    
    def process_folder(self, folder_path):
        results_csv = os.path.join(folder_path, "results.csv")
        images_folder = os.path.join(folder_path, "images")
        
        if not os.path.exists(results_csv):
            print(f"結果ファイルが見つかりません: {results_csv}")
            return
        
        df = pd.read_csv(results_csv)
        
        for _, row in df.iterrows():
            image_file_name = row["image_file_name"]
            image_path = os.path.join(images_folder, image_file_name)
            
            if not os.path.exists(image_path):
                print(f"画像ファイルが見つかりません: {image_path}")
                continue
                
            image = cv2.imread(image_path)
            if image is None:
                print(f"画像の読み込みに失敗しました: {image_path}")
                continue
            
            # 画像を表示
            cv2.imshow('Image', image)
            print(f"現在の画像: {image_file_name}")
            print("Enterキーを押して保存、または新しいpos番号を入力してください（1-99）:")
            
            input_buffer = ""
            while True:
                key = cv2.waitKey(0) & 0xFF
                
                if key == 13:  # Enterキー
                    if input_buffer:  # 数字が入力されている場合
                        new_pos_number = input_buffer
                        new_filename = self.get_new_filename(image_file_name, new_pos_number)
                        
                        while True:
                            output_path = os.path.join(self.output_dir, new_filename)
                            
                            if self.check_duplicate_filename(new_filename):
                                print("\nファイルが重複しています。選択してください：")
                                print("1: 新しいファイル名を入力する")
                                print("2: スキップする")
                                
                                choice = input("選択 (1/2): ").strip()
                                if choice == "1":
                                    print(f"\n現在のファイル名: {new_filename}")
                                    new_name = input("新しいファイル名を入力してください: ").strip()
                                    
                                    if not new_name:
                                        print("入力がキャンセルされました")
                                        break
                                    
                                    if not new_name.endswith(os.path.splitext(image_file_name)[1]):
                                        new_name += os.path.splitext(image_file_name)[1]
                                    
                                    new_filename = new_name
                                    continue
                                    
                                elif choice == "2":
                                    print("スキップします")
                                    break
                                
                                else:
                                    print("無効な選択です。1か2を入力してください。")
                                    continue
                            
                            cv2.imwrite(output_path, image)
                            print(f"保存完了: {output_path}")
                            break
                        break
                        
                    else:  # 直接Enterが押された場合
                        output_path = os.path.join(self.output_dir, image_file_name)
                        if self.check_duplicate_filename(image_file_name):
                            print("\nファイルが重複しています。選択してください：")
                            print("1: 新しいファイル名を入力する")
                            print("2: スキップする")
                            
                            while True:
                                choice = input("選択 (1/2): ").strip()
                                if choice == "1":
                                    while True:
                                        print(f"\n現在のファイル名: {image_file_name}")
                                        new_name = input("新しいファイル名を入力してください: ").strip()
                                        
                                        # 入力が空の場合はスキップ
                                        if not new_name:
                                            print("入力がキャンセルされました")
                                            break
                                        
                                        # 拡張子の確認と追加
                                        if not new_name.endswith(os.path.splitext(image_file_name)[1]):
                                            new_name += os.path.splitext(image_file_name)[1]
                                        
                                        output_path = os.path.join(self.output_dir, new_name)
                                        
                                        # 新しいファイル名でも重複チェック
                                        if self.check_duplicate_filename(new_name):
                                            print("Warning: 指定されたファイル名も既に存在します。")
                                            continue
                                        
                                        # 問題なければ保存して終了
                                        cv2.imwrite(output_path, image)
                                        print(f"保存完了: {output_path}")
                                        break
                                    break
                                
                                elif choice == "2":
                                    print("スキップします")
                                    break
                                
                                else:
                                    print("無効な選択です。1か2を入力してください。")
                            continue
                        cv2.imwrite(output_path, image)
                        print(f"保存完了: {output_path}")
                        break
                
                elif key >= ord('0') and key <= ord('9'):
                    input_buffer += chr(key)
                    print(f"入力中: {input_buffer}")
                    
                elif key == 8:  # バックスペース
                    input_buffer = input_buffer[:-1]
                    print(f"入力中: {input_buffer}")
                
                elif key == 27:  # ESCキー
                    print("処理をスキップします")
                    break
            
            cv2.destroyWindow('Image')

def get_numeric_folders(base_dir="data"):
    """
    dataディレクトリ内の数字名フォルダを取得する
    """
    folders = []
    if os.path.exists(base_dir):
        for item in os.listdir(base_dir):
            if os.path.isdir(os.path.join(base_dir, item)) and item.isdigit():
                folders.append(item)
    return sorted(folders, key=lambda x: int(x))

def main():
    processor = ImageProcessor()
    
    # dataディレクトリ内の数字フォルダを取得
    numeric_folders = get_numeric_folders()
    if not numeric_folders:
        print("処理対象のフォルダが見つかりません")
        return
    
    print("処理するフォルダを選択してください:")
    for i, folder in enumerate(numeric_folders, 1):
        print(f"{i}: {folder}")
    print("0: 全てのフォルダを処理")
    
    while True:
        try:
            choice = input("番号を入力してください: ").strip()
            if not choice:
                continue
                
            choice = int(choice)
            if choice == 0:
                # 全てのフォルダを処理
                folders_to_process = numeric_folders
                break
            elif 1 <= choice <= len(numeric_folders):
                # 選択されたフォルダのみ処理
                folders_to_process = [numeric_folders[choice-1]]
                break
            else:
                print(f"1から{len(numeric_folders)}までの数字、または0を入力してください")
        except ValueError:
            print("有効な数字を入力してください")
    
    # 選択されたフォルダを処理
    for folder in folders_to_process:
        folder_path = os.path.join("data", folder.zfill(2))
        print(f"\n=== フォルダ {folder} の処理を開始 ===")
        processor.process_folder(folder_path)
    
    print("\n処理が完了しました")

if __name__ == "__main__":
    main()