import tkinter as tk
from tkinter import filedialog, messagebox, scrolledtext, ttk
from PIL import Image, ImageTk, ImageGrab
import threading
import os
import numpy as np
import pandas as pd
import cv2
import tempfile
import io
import json
import sys
import os

def resource_path(relative_path):
    """ PyInstaller에 의해 임시 폴더에 생성된 리소스의 절대 경로를 반환합니다. """
    try:
        # PyInstaller는 임시 폴더를 생성하고 _MEIPASS에 경로를 저장합니다.
        base_path = sys._MEIPASS
    except Exception:
        base_path = os.path.abspath(".")

    return os.path.join(base_path, relative_path)

class WD14TaggerGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("WD14 Auto Prompt Generator (aiart)")
        self.root.geometry("700x870")

        self.SETTINGS_FILE = "tagger_app_settings.json"

        self.current_image_path = None
        self.session = None
        self.tags_df = None

        # 변수 초기화
        self.general_threshold_var = tk.DoubleVar(value=0.5)
        self.char_threshold_var = tk.DoubleVar(value=0.85)
        self.rm_c_var = tk.BooleanVar()
        self.rm_color_var = tk.BooleanVar()
        self.rm_clothes_var = tk.BooleanVar()
        self.webui_mode_var = tk.BooleanVar()
        self.instant_inference_var = tk.BooleanVar(value=False)
        ## 1) 접기/펼치기 상태 변수 추가
        self.show_removed_tags_var = tk.BooleanVar(value=True)
        self.show_ignore_tags_var = tk.BooleanVar(value=False)
        self.prompt_active_var = tk.BooleanVar(value=False)


        self.init_filter_data()
        self.setup_ui()
        self.load_model()
        self.show_initial_preview()
        self.setup_drag_drop()

        self.load_settings()
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
        
        ## 2) 전역 Ctrl+V 이벤트 바인딩
        self.root.bind('<Control-v>', self.handle_paste)

    def init_filter_data(self):
        # ... (기존과 동일)
        try:
            from tagbag import bag_of_tags, clothes_list
            from character_dictionary import character_dictionary as cd

            self.bag_of_tags = bag_of_tags
            self.clothes_list = clothes_list
            self.character_keys = list(cd.keys()) if isinstance(cd, dict) else []

        except ImportError as e:
            print(f"⚠ NAIA 파일 import 실패: {e}")
            print("기본 데이터를 사용합니다.")
            self.bag_of_tags = [
                'smile', 'blush', 'open_mouth', 'closed_eyes', 'wink', 'frown',
                'twintails', 'ponytail', 'braid', 'long_hair', 'short_hair',
                'large_breasts', 'small_breasts', 'cleavage', 'navel'
            ]
            self.clothes_list = [
                'dress', 'skirt', 'shirt', 'jacket', 'uniform', 'school_uniform'
            ]
            self.character_keys = []

        self.colors = ['black', 'white', 'blond', 'silver', 'gray', 'yellow',
                        'blue', 'purple', 'red', 'pink', 'brown', 'orange',
                        'green', 'aqua', 'gradient']
    
    ## 1) 접기/펼치기 토글 함수
    def toggle_removed_tags(self):
        if self.show_removed_tags_var.get():
            self.removed_tags_text.grid(row=3, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        else:
            self.removed_tags_text.grid_forget()

    ## 1) 접기/펼치기 토글 함수
    def toggle_ignore_tags(self):
        if self.show_ignore_tags_var.get():
            self.ignore_tags_text.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S))
            self.save_settings_btn.grid(row=1, column=2, sticky="ne", padx=(10, 0))
        else:
            self.ignore_tags_text.grid_forget()
            self.save_settings_btn.grid_forget()

    def toggle_prompts(self):
        pass

    def setup_ui(self):
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # ... (상단 UI 설정은 기존과 거의 동일)
        load_frame = ttk.Frame(main_frame)
        load_frame.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(0, 10))
        self.load_btn = ttk.Button(load_frame, text="이미지 불러오기", command=self.load_image)
        self.load_btn.grid(row=0, column=0, padx=(0, 5))
        self.clipboard_btn = ttk.Button(load_frame, text="이미지 불러오기 (클립보드)", command=self.load_from_clipboard)
        self.clipboard_btn.grid(row=0, column=1)

        content_frame = ttk.Frame(main_frame)
        content_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), pady=(0, 10))

        self.preview_frame = ttk.LabelFrame(content_frame, text="미리보기", padding="10")
        self.preview_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S), padx=(0, 2))
        self.image_label = ttk.Label(self.preview_frame, text="")
        self.image_label.grid(row=0, column=0)

        settings_frame = ttk.LabelFrame(content_frame, text="설정", padding="10")
        settings_frame.grid(row=0, column=1, sticky=(tk.W, tk.E, tk.N, tk.S), padx=(2, 0))

        ttk.Label(settings_frame, text="태그 임계값:").grid(row=0, column=0, sticky=tk.W)
        ttk.Spinbox(settings_frame, from_=0.0, to=1.0, increment=0.05,
                   textvariable=self.general_threshold_var, width=15).grid(row=0, column=1, pady=(0, 10))

        ttk.Label(settings_frame, text="캐릭터 임계값:").grid(row=1, column=0, sticky=tk.W)
        ttk.Spinbox(settings_frame, from_=0.0, to=1.0, increment=0.05,
                   textvariable=self.char_threshold_var, width=15).grid(row=1, column=1, pady=(0, 10))
        
        extract_control_frame = ttk.Frame(settings_frame)
        extract_control_frame.grid(row=2, column=0, columnspan=2, pady=10)

        self.extract_btn = ttk.Button(extract_control_frame, text="태그 추출",
                                    command=self.extract_tags, state='disabled')
        self.extract_btn.pack(side=tk.LEFT, padx=(0, 10))

        self.instant_inference_cb = ttk.Checkbutton(extract_control_frame, text="즉시 추론", variable=self.instant_inference_var)
        self.instant_inference_cb.pack(side=tk.LEFT)

        filter_frame = ttk.LabelFrame(settings_frame, text="태그 필터링", padding="5")
        filter_frame.grid(row=3, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(10, 0))
        ttk.Checkbutton(filter_frame, text="캐릭터 특징+이름 제거", variable=self.rm_c_var).grid(row=0, column=0, sticky=tk.W, pady=2)
        ttk.Checkbutton(filter_frame, text="의류 색상 제거", variable=self.rm_color_var).grid(row=1, column=0, sticky=tk.W, pady=2)
        ttk.Checkbutton(filter_frame, text="의상 제거 (비추천)", variable=self.rm_clothes_var).grid(row=2, column=0, sticky=tk.W, pady=2)
        ttk.Checkbutton(filter_frame, text="WEBUI/Comfy 모드", variable=self.webui_mode_var).grid(row=3, column=0, sticky=tk.W, pady=2)

        # 선행/후행 프롬프트 프레임
        prompt_frame = ttk.Frame(settings_frame)
        prompt_frame.grid(row=4, column=0, columnspan=2, sticky="ew", pady=(15, 0))
        
        self.prompt_active_cb = ttk.Checkbutton(prompt_frame, text="선행/후행 프롬프트 활성",
                                                variable=self.prompt_active_var,
                                                command=self.toggle_prompts)
        self.prompt_active_cb.grid(row=0, column=0, sticky="w")

        # 탭 위젯 (Notebook) 생성
        self.prompt_notebook = ttk.Notebook(prompt_frame)
        self.prompt_notebook.grid(row=1, column=0, columnspan=2, sticky="ew", pady=(5, 0))

        # 선행 프롬프트 탭
        prefix_frame = ttk.Frame(self.prompt_notebook, padding=5)
        self.prompt_notebook.add(prefix_frame, text="선행 프롬프트")
        self.prefix_text = scrolledtext.ScrolledText(prefix_frame, height=3, width=30)
        self.prefix_text.pack(fill="both", expand=True)

        # 후행 프롬프트 탭
        suffix_frame = ttk.Frame(self.prompt_notebook, padding=5)
        self.prompt_notebook.add(suffix_frame, text="후행 프롬프트")
        self.suffix_text = scrolledtext.ScrolledText(suffix_frame, height=3, width=30)
        self.suffix_text.pack(fill="both", expand=True)
        
        # 초기 상태 설정
        self.toggle_prompts()

        result_frame = ttk.LabelFrame(main_frame, text="추출된 태그", padding="10")
        result_frame.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), pady=(10, 0))

        self.result_text = scrolledtext.ScrolledText(result_frame, height=8)
        self.result_text.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        self.result_text.tag_config("prompt", foreground="grey")

        copy_frame = ttk.Frame(result_frame)
        copy_frame.grid(row=1, column=0, pady=(10, 5), sticky=(tk.W, tk.E))
        ttk.Button(copy_frame, text="전체 복사", command=self.copy_all_tags).grid(row=0, column=0, padx=(0, 5))
        ttk.Button(copy_frame, text="2번 태그부터 복사", command=self.copy_from_second).grid(row=0, column=1, padx=5)
        ttk.Button(copy_frame, text="3번 태그부터 복사", command=self.copy_from_third).grid(row=0, column=2, padx=5)

        ## 1) '제거된 태그' 라벨을 체크박스로 변경
        self.removed_tags_cb = ttk.Checkbutton(result_frame, text="제거된 태그",
                                               variable=self.show_removed_tags_var,
                                               command=self.toggle_removed_tags)
        self.removed_tags_cb.grid(row=2, column=0, sticky=tk.W, pady=(10, 2))
        
        self.removed_tags_text = scrolledtext.ScrolledText(result_frame, height=5)
        # self.removed_tags_text.grid(...)는 toggle 함수에서 관리

        ## 1) '무시할 태그' 라벨을 체크박스로 변경하고 별도 프레임으로 분리
        ignore_frame = ttk.Frame(main_frame, padding="0")
        ignore_frame.grid(row=3, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=(10, 0))
        
        self.ignore_tags_cb = ttk.Checkbutton(ignore_frame, text="무시할 태그 (쉼표로 구분)",
                                              variable=self.show_ignore_tags_var,
                                              command=self.toggle_ignore_tags)
        self.ignore_tags_cb.grid(row=0, column=0, sticky=tk.W)

        self.ignore_tags_text = scrolledtext.ScrolledText(ignore_frame, height=4)
        self.save_settings_btn = ttk.Button(ignore_frame, text="설정 저장", command=self.save_settings)

        ignore_frame.columnconfigure(1, weight=1)

        ## 1) 초기 상태 설정
        self.toggle_removed_tags()
        self.toggle_ignore_tags()

        # ... (그리드 설정)
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(0, weight=1)
        # main_frame.rowconfigure(2)와 (3)은 내용에 따라 자동 조절되도록 weight 미설정
        content_frame.columnconfigure(0, weight=1)
        content_frame.columnconfigure(1, weight=1)
        result_frame.columnconfigure(0, weight=1)
        result_frame.rowconfigure(0, weight=1)


    def on_closing(self):
        self.save_settings()
        self.root.destroy()

    def save_settings(self):
        settings = {
            'general_threshold': self.general_threshold_var.get(),
            'char_threshold': self.char_threshold_var.get(),
            'rm_c': self.rm_c_var.get(),
            'rm_color': self.rm_color_var.get(),
            'rm_clothes': self.rm_clothes_var.get(),
            'webui_mode': self.webui_mode_var.get(),
            'instant_inference': self.instant_inference_var.get(),
            'ignore_tags': self.ignore_tags_text.get(1.0, tk.END).strip(),
            ## 1) 접기/펼치기 상태 저장
            'show_removed_tags': self.show_removed_tags_var.get(),
            'show_ignore_tags': self.show_ignore_tags_var.get(),
            'prompt_active': self.prompt_active_var.get(),
            'prefix_prompt': self.prefix_text.get(1.0, tk.END).strip(),
            'suffix_prompt': self.suffix_text.get(1.0, tk.END).strip()
        }
        try:
            with open(self.SETTINGS_FILE, 'w', encoding='utf-8') as f:
                json.dump(settings, f, indent=4)
            # if self.root.focus_get() == self.save_settings_btn:
            #     messagebox.showinfo("저장 완료", f"설정이 {self.SETTINGS_FILE}에 저장되었습니다.")
        except Exception as e:
            messagebox.showerror("저장 실패", f"설정 저장 중 오류가 발생했습니다: {e}")

    def load_settings(self):
        try:
            if os.path.exists(self.SETTINGS_FILE):
                with open(self.SETTINGS_FILE, 'r', encoding='utf-8') as f:
                    settings = json.load(f)
                    # ... (기존 설정 로드)
                    self.general_threshold_var.set(settings.get('general_threshold', 0.5))
                    self.char_threshold_var.set(settings.get('char_threshold', 0.85))
                    self.rm_c_var.set(settings.get('rm_c', False))
                    self.rm_color_var.set(settings.get('rm_color', False))
                    self.rm_clothes_var.set(settings.get('rm_clothes', False))
                    self.webui_mode_var.set(settings.get('webui_mode', False))
                    self.instant_inference_var.set(settings.get('instant_inference', False))
                    
                    self.ignore_tags_text.delete(1.0, tk.END)
                    self.ignore_tags_text.insert(1.0, settings.get('ignore_tags', ''))

                    ## 1) 접기/펼치기 상태 로드 및 UI 업데이트
                    self.show_removed_tags_var.set(settings.get('show_removed_tags', True))
                    self.show_ignore_tags_var.set(settings.get('show_ignore_tags', False))
                    self.toggle_removed_tags()
                    self.toggle_ignore_tags()
                    self.prompt_active_var.set(settings.get('prompt_active', False))
                    
                    self.prefix_text.delete(1.0, tk.END)
                    self.prefix_text.insert(1.0, settings.get('prefix_prompt', ''))
                    
                    self.suffix_text.delete(1.0, tk.END)
                    self.suffix_text.insert(1.0, settings.get('suffix_prompt', ''))

                    self.toggle_prompts() # UI 상태 업데이트

        except Exception as e:
            messagebox.showwarning("설정 로드 실패", f"설정 파일({self.SETTINGS_FILE})을 불러오는 데 실패했습니다: {e}")
    
    ## 2) 전역 붙여넣기 이벤트 핸들러 추가
    def handle_paste(self, event):
        try:
            # 1. 클립보드에 이미지가 있는지 확인
            clipboard_image = ImageGrab.grabclipboard()
            if isinstance(clipboard_image, Image.Image):
                self.load_from_clipboard()
                return "break"  # 이벤트 전파 중단
        except (ValueError, TypeError):
             pass
        except Exception as e:
            print(f"클립보드 이미지 확인 중 오류: {e}")


        try:
            # 클립보드의 텍스트 내용을 가져옵니다.
            clipboard_text = self.root.clipboard_get()
            
            # 2. 클립보드에 텍스트가 있고, URL 형식인지 확인
            if isinstance(clipboard_text, str) and clipboard_text.startswith(('http://', 'https://')):
                self._load_image_from_url(clipboard_text)
                return "break" # 이벤트 전파 중단
            
            ## --- 여기부터 추가되는 부분 --- ##
            # 3. 클립보드 텍스트가 이미지 파일 경로인지 확인
            image_extensions = ('.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tiff')
            if isinstance(clipboard_text, str) and clipboard_text.lower().endswith(image_extensions):
                # 따옴표로 감싸진 경로일 경우 제거 (파일 탐색기에서 복사 시)
                path = clipboard_text.strip('"')
                # 파일 시스템에 실제 존재하는 경로인지 확인 후 헬퍼 메서드 호출
                if os.path.exists(path):
                    self._load_image_from_path(path)
                    return "break" # 이벤트 전파 중단
            ## --- 여기까지 추가 --- ##

        except tk.TclError:
            # 클립보드에 텍스트가 없는 경우
            pass
        except Exception as e:
            messagebox.showerror("붙여넣기 오류", f"처리 중 오류 발생: {e}")
            return "break"

        # 위 조건에 해당하지 않으면 기본 붙여넣기 동작
        return

    ## 2) URL 처리 로직을 별도 메서드로 분리 (재사용성)
    def _load_image_from_url(self, url):
        try:
            img = self.download_image_from_url(url)
            temp_dir = tempfile.gettempdir()
            temp_file = os.path.join(temp_dir, f"wd14_url_{os.getpid()}.png")
            img.save(temp_file)
            
            self.current_image_path = temp_file
            self.show_preview(temp_file)
            if self.session is not None:
                self.extract_btn.config(state='normal')
                if self.instant_inference_var.get():
                    self.extract_tags()
        except Exception as e:
            messagebox.showerror("오류", f"웹 이미지 로드 실패: {str(e)}")

    def load_model(self):
        try:
            import onnxruntime as ort
            
            # --- 수정 후 ---
            model_path = resource_path("model.onnx")
            tags_path = resource_path("selected_tags.csv")
            
            self.session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
            self.tags_df = pd.read_csv(tags_path)
            
        except Exception as e:
            messagebox.showerror("오류", f"모델 로드 실패: {str(e)}")

    def show_initial_preview(self):
        # ... (기존과 동일)
        try:
            white_image = Image.new('RGB', (374, 374), 'white')
            photo = ImageTk.PhotoImage(white_image)
            self.image_label.config(image=photo, text="이미지를 Drag&Drop\n해주세요")
            self.image_label.image = photo
        except Exception as e:
            self.image_label.config(text="이미지를 Drag&Drop\n해주세요")

    def setup_drag_drop(self):
        # ... (기존과 동일)
        try:
            from tkinterdnd2 import TkinterDnD, DND_ALL
            self.TkdndVersion = TkinterDnD._require(self.root)
            self.image_label.drop_target_register(DND_ALL)
            self.image_label.dnd_bind('<<Drop>>', self.on_drop)
            self.image_label.dnd_bind('<<DragEnter>>', self.on_drag_enter)
            self.image_label.dnd_bind('<<DragLeave>>', self.on_drag_leave)
        except ImportError:
            print("tkinterdnd2가 설치되지 않음. pip install tkinterdnd2로 설치하면 드래그앤드롭을 사용할 수 있습니다.")
        except Exception as e:
            print(f"Drag & Drop 초기화 오류: {e}")

    # ... (on_drag_enter, on_drag_leave, download_image_from_url 기존과 동일)
    def on_drag_enter(self, event):
        try:
            self.preview_frame.config(text="미리보기 (이미지를 놓아주세요!)")
        except: pass

    def on_drag_leave(self, event):
        try:
            self.preview_frame.config(text="미리보기")
        except: pass

    def download_image_from_url(self, url):
        import requests
        response = requests.get(url)
        response.raise_for_status()
        return Image.open(io.BytesIO(response.content))

    def on_drop(self, event):
        try:
            if not isinstance(event, str):
                file_path_or_url = event.data.strip("{}")
            else:
                file_path_or_url = event

            if not file_path_or_url:
                messagebox.showwarning("경고", "드롭된 데이터를 읽을 수 없습니다.")
                self.preview_frame.config(text="미리보기")
                return
            
            # URL 처리
            if file_path_or_url.startswith(('http://', 'https://')):
                self._load_image_from_url(file_path_or_url) ## 2) 재사용되는 URL 로드 함수 호출
                self.preview_frame.config(text="미리보기")
                return
            else:
                if file_path_or_url.startswith(('blob:')):
                    messagebox.showwarning("Blob URL 오류", "NAI 홈페이지에서 생성한 이미지는 즉시 복사할 수 없습니다.\n이미지를 다운로드한 후 드래그 & 드롭해주세요.")
                    self.preview_frame.config(text="미리보기")
                    return
                
                # 파일 처리
                file_path_or_url = file_path_or_url.replace("\\", '/')
                if os.path.exists(file_path_or_url):
                    self.current_image_path = file_path_or_url
                    self.show_preview(file_path_or_url)
                    if self.session is not None:
                        self.extract_btn.config(state='normal')
                        if self.instant_inference_var.get():
                            self.extract_tags()
                    self.preview_frame.config(text="미리보기")
                else:
                    messagebox.showerror("오류", f"유효하지 않은 파일 경로입니다: '{file_path_or_url}'")
                    self.preview_frame.config(text="미리보기")
                    
        except Exception as e:
            messagebox.showerror("오류", f"파일 드롭 오류: {str(e)}")
            self.preview_frame.config(text="미리보기")
    
    def _load_image_from_path(self, path):
        """주어진 파일 경로에서 이미지를 로드하고 미리보기를 표시합니다."""
        # 경로의 유효성을 한 번 더 확인합니다.
        if os.path.exists(path):
            self.current_image_path = path
            self.show_preview(path)
            if self.session is not None:
                self.extract_btn.config(state='normal')
                if self.instant_inference_var.get():
                    self.extract_tags()
        else:
            # 혹시 모를 경우를 대비한 경고창
            messagebox.showwarning("경로 오류", f"파일을 찾을 수 없습니다: {path}")

    def load_from_clipboard(self):
        try:
            image = ImageGrab.grabclipboard()
            if isinstance(image, Image.Image):
                temp_dir = tempfile.gettempdir()
                temp_file = os.path.join(temp_dir, f"wd14_clipboard_{os.getpid()}.png")
                # RGBA 이미지를 RGB로 변환하여 JPEG/PNG 호환성 문제 방지
                if image.mode == 'RGBA':
                    image = image.convert('RGB')
                image.save(temp_file, 'PNG') # 포맷을 명시적으로 지정
                
                self.current_image_path = temp_file
                self.show_preview(temp_file)
                if self.session is not None:
                    self.extract_btn.config(state='normal')
                    if self.instant_inference_var.get():
                        self.extract_tags()
            else:
                # 이 메시지는 handle_paste 로직 때문에 거의 보이지 않게 됨
                messagebox.showinfo("정보", "클립보드에 이미지가 없습니다.")
        except Exception as e:
            messagebox.showerror("오류", f"클립보드 불러오기 오류: {str(e)}")

    # ... (load_image, show_preview, extract_tags, _extract_tags_thread, apply_tag_filters,
    #  _is_character_name, _format_results, copy 함수들, _update_results, make_square, smart_resize 기존과 동일)
    def load_image(self):
        file_path = filedialog.askopenfilename(
            title="이미지 선택",
            filetypes=[("이미지 파일", "*.jpg *.jpeg *.png *.bmp *.tiff *.webp")]
        )
        if file_path:
            self.current_image_path = file_path
            self.show_preview(file_path)
            if self.session is not None:
                self.extract_btn.config(state='normal')
                if self.instant_inference_var.get():
                    self.extract_tags()

    def show_preview(self, image_path):
        try:
            image = Image.open(image_path).convert('RGB')
            size = max(image.size)
            square_image = Image.new('RGB', (size, size), 'white')
            paste_x = (size - image.width) // 2
            paste_y = (size - image.height) // 2
            square_image.paste(image, (paste_x, paste_y))
            square_image.thumbnail((374, 374), Image.Resampling.LANCZOS)
            photo = ImageTk.PhotoImage(square_image)
            self.image_label.config(image=photo, text="")
            self.image_label.image = photo
        except Exception as e:
            self.image_label.config(text=f"이미지 로드 실패: {str(e)}")

    def extract_tags(self):
        if not self.current_image_path or not self.session:
            return
        self.extract_btn.config(state='disabled', text='추출 중...')
        self.result_text.delete(1.0, tk.END)
        self.result_text.insert(tk.END, "태그 추출 중...")
        self.removed_tags_text.delete(1.0, tk.END) 
        thread = threading.Thread(target=self._extract_tags_thread)
        thread.daemon = True
        thread.start()

    def _extract_tags_thread(self):
        try:
            input_image = Image.open(self.current_image_path)
            _, height, _, _ = self.session.get_inputs()[0].shape
            image = input_image.convert('RGBA')
            new_image = Image.new('RGBA', image.size, 'WHITE')
            new_image.paste(image, mask=image)
            image = new_image.convert('RGB')
            image = np.asarray(image)[:, :, ::-1]
            image = self.make_square(image, height)
            image = self.smart_resize(image, height)
            image = image.astype(np.float32)
            image = np.expand_dims(image, 0)
            input_name = self.session.get_inputs()[0].name
            label_name = self.session.get_outputs()[0].name
            confidents = self.session.run([label_name], {input_name: image})[0]
            tags_with_conf = []
            for i, conf in enumerate(confidents[0]):
                if i < len(self.tags_df):
                    tags_with_conf.append((self.tags_df.iloc[i]['name'], float(conf)))
            tags = dict(tags_with_conf[4:])
            
            result_tuple = self._format_results(tags)
            self.root.after(0, self._update_results, result_tuple)
        except Exception as e:
            error_msg = f"태그 추출 오류: {str(e)}"
            self.root.after(0, self._update_results, (error_msg, ""))

    def apply_tag_filters(self, tags_list):
        """태그 필터링 적용 (공백 기준으로 로직 단순화)"""
        # 모델에서 넘어온 원본 태그(_ 포함) 리스트
        original_tags = tags_list.copy()
        # 모든 처리를 위해 공백 기준으로 변환한 태그 리스트
        tags_to_process = [tag.replace('_', ' ') for tag in original_tags]
        
        # 최종적으로 남을 태그들의 인덱스를 관리
        final_indices = list(range(len(tags_to_process)))
        
        user_removed_originals = []
        other_removed_originals = []

        # 1. 사용자가 정의한 무시할 태그 처리
        ignore_list_str = self.ignore_tags_text.get(1.0, tk.END).strip()
        if ignore_list_str:
            user_ignored_tags = {tag.strip().replace('_', ' ') for tag in ignore_list_str.split(',') if tag.strip()}
            indices_to_remove = {i for i in final_indices if tags_to_process[i] in user_ignored_tags}
            for i in indices_to_remove:
                user_removed_originals.append(original_tags[i])
            final_indices = [i for i in final_indices if i not in indices_to_remove]

        # 2. 나머지 필터들 적용 (내부 데이터가 이미 공백 기준이므로 직접 비교)
        if self.rm_c_var.get():
            indices_to_remove = {i for i in final_indices if tags_to_process[i] in self.bag_of_tags or self._is_character_name(tags_to_process[i])}
            for i in indices_to_remove:
                other_removed_originals.append(original_tags[i])
            final_indices = [i for i in final_indices if i not in indices_to_remove]

        if self.rm_color_var.get():
            indices_to_remove = {i for i in final_indices if any(color in tags_to_process[i] for color in self.colors) and " eyes" not in tags_to_process[i] and " hair" not in tags_to_process[i] and " pupils" not in tags_to_process[i]}
            for i in indices_to_remove:
                other_removed_originals.append(original_tags[i])
            final_indices = [i for i in final_indices if i not in indices_to_remove]
        
        if self.rm_clothes_var.get():
            indices_to_remove = {i for i in final_indices if tags_to_process[i] in self.clothes_list and tags_to_process[i] not in self.bag_of_tags}
            for i in indices_to_remove:
                other_removed_originals.append(original_tags[i])
            final_indices = [i for i in final_indices if i not in indices_to_remove]

        # 3. 최종 태그 목록 생성 및 포맷팅
        final_tags_with_space = [tags_to_process[i] for i in final_indices]
        
        if self.webui_mode_var.get():
            final_tags_formatted = [tag.replace('(', '\\(').replace(')', '\\)') for tag in final_tags_with_space]
        else:
            final_tags_formatted = final_tags_with_space

        # 4. 선행/후행 프롬프트 처리
        final_tags_with_type = [(tag, 'original') for tag in final_tags_formatted]
        if self.prompt_active_var.get():
            prefix_str = self.prefix_text.get(1.0, tk.END).strip()
            if prefix_str:
                prefix_tags = [t.strip() for t in prefix_str.split(',') if t.strip()]
                prefix_tags_with_type = [(tag, 'prompt') for tag in prefix_tags]
                person_tags = {'1girl', '2girls', '3girls', '4girls', '5girls', '6+girls', 
                            '1boy', '2boys', '3boys', '4boys', '5boys', '6+boys'}
                last_person_tag_index = -1
                search_range = min(len(final_tags_with_type), 4)
                for i in range(search_range):
                    if final_tags_with_type[i][0] in person_tags:
                        last_person_tag_index = i
                if last_person_tag_index != -1:
                    final_tags_with_type[last_person_tag_index+1:last_person_tag_index+1] = prefix_tags_with_type
                else:
                    final_tags_with_type[0:0] = prefix_tags_with_type
            
            suffix_str = self.suffix_text.get(1.0, tk.END).strip()
            if suffix_str:
                suffix_tags = [t.strip() for t in suffix_str.split(',') if t.strip()]
                suffix_tags_with_type = [(tag, 'prompt') for tag in suffix_tags]
                final_tags_with_type.extend(suffix_tags_with_type)

        return final_tags_with_type, user_removed_originals, list(set(other_removed_originals))

    def _is_character_name(self, tag):
        if hasattr(self, 'character_keys') and tag in self.character_keys:
            return True
        if ('_' in tag and any(c.isupper() for c in tag) and
            not tag.startswith('1') and
            not any(word in tag.lower() for word in ['girl', 'boy', 'solo', 'multiple']) and
            not any(word in tag.lower() for word in ['hair', 'eye', 'dress', 'shirt']) and
            len(tag) > 3):
            return True
        return False

    def _format_results(self, tags):
        """결과 포맷팅 (타입 정보가 포함된 리스트를 그대로 전달)"""
        # ... (메서드 상단의 1차 필터링 로직은 동일) ...
        general_threshold = self.general_threshold_var.get()
        char_threshold = self.char_threshold_var.get()
        
        person_tags = {'1girl', '2girls', '3girls', '4girls', '5girls', '6+girls',
                       '1boy', '2boys', '3boys', '4boys', '5boys', '6+boys',
                       'multiple_girls', 'multiple_boys', 'solo', 'no_humans'}
        
        filtered_tags = []
        for tag, conf in tags.items():
            if tag in person_tags:
                if conf >= 0.9: filtered_tags.append(tag)
            elif any(c.isupper() for c in tag) and '_' in tag and conf >= char_threshold:
                filtered_tags.append(tag)
            elif conf >= general_threshold:
                filtered_tags.append(tag)

        # 2차 필터링: 타입 정보가 포함된 최종 태그 리스트를 받음
        final_tags_with_type, user_removed, other_removed = self.apply_tag_filters(filtered_tags)
        
        # 제거된 태그들을 표시용으로 재구성
        removed_for_display = []
        for tag in user_removed:
            removed_for_display.append((tag, 'ignored'))
        for tag in other_removed:
            if tag not in user_removed:
                removed_for_display.append((tag, 'other'))
        
        # 튜플로 결과 반환: ((태그, 타입) 리스트, (태그, 타입) 리스트)
        return final_tags_with_type, removed_for_display

    def copy_all_tags(self):
        content = self.result_text.get(1.0, tk.END).strip()
        if content:
            self.root.clipboard_clear()
            self.root.clipboard_append(content)
            #messagebox.showinfo("복사 완료", "전체 태그가 클립보드에 복사되었습니다.")
        else:
            messagebox.showwarning("복사 실패", "복사할 태그가 없습니다.")

    def copy_from_second(self):
        self._copy_from_index(1, "2번 태그부터")

    def copy_from_third(self):
        self._copy_from_index(2, "3번 태그부터")

    def _copy_from_index(self, start_index, description):
        content = self.result_text.get(1.0, tk.END).strip()
        if content:
            tags_list = [tag.strip() for tag in content.split(',')]
            if len(tags_list) > start_index:
                selected_tags = tags_list[start_index:]
                result = ', '.join(selected_tags)
                self.root.clipboard_clear()
                self.root.clipboard_append(result)
                #messagebox.showinfo("복사 완료", f"{description} 태그가 클립보드에 복사되었습니다.")
            else:
                messagebox.showwarning("복사 실패", f"태그가 {start_index + 1}개 미만입니다.")
        else:
            messagebox.showwarning("복사 실패", "복사할 태그가 없습니다.")

    def _update_results(self, result_tuple):
        """결과 업데이트 (선행/후행 프롬프트 색상 적용)"""
        main_tags_info, removed_tags_info = result_tuple
        
        # 메인 태그 박스 업데이트
        self.result_text.delete(1.0, tk.END)
        if main_tags_info:
            for i, (tag, tag_type) in enumerate(main_tags_info):
                if i > 0:
                    self.result_text.insert(tk.END, ', ')
                
                # 프롬프트 태그일 경우 "prompt" 스타일 적용
                if tag_type == 'prompt':
                    self.result_text.insert(tk.END, tag, 'prompt')
                else: # 'original' 태그
                    self.result_text.insert(tk.END, tag)
        
        # 제거된 태그 박스 업데이트 (기존과 동일)
        self.removed_tags_text.config(state=tk.NORMAL)
        self.removed_tags_text.delete(1.0, tk.END)
        if removed_tags_info:
            for i, (tag, tag_type) in enumerate(removed_tags_info):
                display_tag = tag.replace('_', ' ')
                
                if i > 0:
                    self.removed_tags_text.insert(tk.END, ', ')
                
                if tag_type == 'ignored':
                    self.removed_tags_text.insert(tk.END, display_tag, 'ignored')
                else:
                    self.removed_tags_text.insert(tk.END, display_tag)
        self.removed_tags_text.config(state=tk.DISABLED)

        self.extract_btn.config(state='normal', text='태그 추출')

    def make_square(self, img, target_size):
        old_size = img.shape[:2]
        desired_size = max(old_size)
        desired_size = max(desired_size, target_size)
        delta_w = desired_size - old_size[1]
        delta_h = desired_size - old_size[0]
        top, bottom = delta_h // 2, delta_h - (delta_h // 2)
        left, right = delta_w // 2, delta_w - (delta_w // 2)
        color = [255, 255, 255]
        new_im = cv2.copyMakeBorder(
            img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color
        )
        return new_im

    def smart_resize(self, img, size):
        if img.shape[0] > size:
            img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
        elif img.shape[0] < size:
            img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC)
        return img

def main():
    try:
        from tkinterdnd2 import TkinterDnD
        root = TkinterDnD.Tk()
    except ImportError:
        print("tkinterdnd2가 설치되지 않아 Drag & Drop 기능이 비활성화됩니다.")
        root = tk.Tk()

    app = WD14TaggerGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()