#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
YOLO Annotation Utility Tool
YOLO format annotation file manipulation utility

Usage examples:
    # Interactive mode (recommended)
    python yolo_annotation_modify.py

    # Remove classes
    python yolo_annotation_modify.py remove --dataset dataset/yolo --classes "0,1"

    # Merge classes
    python yolo_annotation_modify.py merge --dataset dataset/yolo --source "0,1,2" --target "vehicle"

    # Rename classes
    python yolo_annotation_modify.py rename --dataset dataset/yolo --mapping "0:vehicle,1:person"

    # Show info
    python yolo_annotation_modify.py info --dataset dataset/yolo
"""

import yaml
import argparse
from pathlib import Path
from typing import Dict, List, Set, Union
import shutil
from datetime import datetime
import sys
import os
from collections import defaultdict


class YOLOAnnotationEditor:
    """Class for editing YOLO format annotations"""

    def __init__(self, dataset_path: str, dry_run: bool = False):
        """
        Args:
            dataset_path: Path to YOLO dataset directory (containing data.yaml)
            dry_run: If True, only simulate changes without writing files
        """
        self.dataset_path = Path(dataset_path)
        self.dry_run = dry_run

        if self.dry_run:
            print("[DRY RUN] 시뮬레이션 모드 - 파일은 수정되지 않습니다\n")

        # Validate dataset path exists
        if not self.dataset_path.exists():
            raise FileNotFoundError(f"Dataset directory not found: {dataset_path}")

        if not self.dataset_path.is_dir():
            raise NotADirectoryError(f"Path is not a directory: {dataset_path}")

        self.yaml_path = self.dataset_path / "data.yaml"

        if not self.yaml_path.exists():
            raise FileNotFoundError(f"data.yaml not found in {dataset_path}")

        self.config = self._load_yaml()
        self.labels_path = self.dataset_path / "labels"

        # Validate labels directory exists
        if not self.labels_path.exists():
            raise FileNotFoundError(f"labels directory not found in {dataset_path}")

        # Find all label files
        self.label_files = list(self.labels_path.rglob("*.txt"))
        print(f"[+] Found {len(self.label_files)} label files")

    def _load_yaml(self) -> dict:
        """Load YOLO data.yaml file"""
        try:
            with open(self.yaml_path, 'r', encoding='utf-8') as f:
                data = yaml.safe_load(f)

            # Validate YAML structure
            if not isinstance(data, dict):
                raise ValueError(f"Invalid YAML format: Expected dictionary, got {type(data)}")

            if 'names' not in data:
                raise ValueError("data.yaml must contain 'names' field")

            if not isinstance(data['names'], dict):
                raise ValueError("'names' field must be a dictionary")

            return data

        except yaml.YAMLError as e:
            raise ValueError(f"Failed to parse YAML file: {e}")
        except Exception as e:
            raise RuntimeError(f"Error loading data.yaml: {e}")

    def backup_labels(self) -> Path:
        """
        Backup entire labels folder

        Returns:
            Path to backup folder
        """
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        backup_path = self.dataset_path / f"labels_backup_{timestamp}"

        if self.dry_run:
            print(f"[DRY RUN] labels 폴더 백업이 생성될 예정입니다: {backup_path}")
            return backup_path

        try:
            print(f"[*] labels 폴더 백업 중...")
            shutil.copytree(self.labels_path, backup_path)
            print(f"[+] 백업 생성됨: {backup_path}")
            return backup_path
        except PermissionError:
            raise PermissionError(f"Permission denied creating backup at {backup_path}")
        except OSError as e:
            raise OSError(f"Failed to create backup: {e}")

    def save_yaml(self, backup: bool = True):
        """
        Save modified data.yaml

        Args:
            backup: Whether to backup original file
        """
        if self.dry_run:
            print(f"[DRY RUN] data.yaml이 저장될 예정입니다: {self.yaml_path}")
            if backup:
                print(f"[DRY RUN] 백업이 생성될 예정입니다")
            return

        try:
            if backup and self.yaml_path.exists():
                backup_path = self.yaml_path.parent / f"data_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.yaml"
                shutil.copy2(self.yaml_path, backup_path)
                print(f"[+] data.yaml 백업 생성됨: {backup_path}")

            with open(self.yaml_path, 'w', encoding='utf-8') as f:
                yaml.dump(self.config, f, allow_unicode=True, sort_keys=False)
            print(f"[+] data.yaml 저장 완료")

        except PermissionError:
            raise PermissionError(f"Permission denied writing to {self.yaml_path}")
        except Exception as e:
            raise RuntimeError(f"Failed to save data.yaml: {e}")

    def get_class_info(self) -> Dict:
        """Get class information and count annotations"""
        names = self.config.get('names', {})

        # Count annotations per class
        class_counts = defaultdict(int)
        total_annotations = 0

        for label_file in self.label_files:
            try:
                with open(label_file, 'r', encoding='utf-8') as f:
                    for line_num, line in enumerate(f, 1):
                        line = line.strip()
                        if line:
                            parts = line.split()
                            if len(parts) < 5:
                                print(f"[!] Warning: Invalid annotation format in {label_file}:{line_num} (expected 5 values, got {len(parts)})")
                                continue
                            try:
                                class_id = int(parts[0])
                                class_counts[class_id] += 1
                                total_annotations += 1
                            except ValueError:
                                print(f"[!] Warning: Invalid class ID in {label_file}:{line_num}")
            except UnicodeDecodeError:
                print(f"[!] Error: Cannot decode file {label_file} (invalid encoding)")
            except PermissionError:
                print(f"[!] Error: Permission denied reading {label_file}")
            except Exception as e:
                print(f"[!] Error reading {label_file}: {e}")

        info = {}
        for class_id, class_name in names.items():
            info[class_id] = {
                'id': class_id,
                'name': class_name,
                'annotation_count': class_counts.get(class_id, 0)
            }

        return info

    def print_info(self):
        """Print annotation information"""
        print("\n" + "="*60)
        print(f"Dataset: {self.dataset_path}")
        print("="*60)

        names = self.config.get('names', {})
        print(f"\n전체 통계:")
        print(f"  - 라벨 파일 수: {len(self.label_files)}")
        print(f"  - 클래스 수: {len(names)}")

        print(f"\n클래스 상세:")
        print(f"{'ID':<6} {'클래스명':<35} {'어노테이션 수':<15}")
        print("-" * 60)

        class_info = self.get_class_info()
        total_annotations = 0
        for class_id in sorted(class_info.keys()):
            info = class_info[class_id]
            ann_count = info['annotation_count']
            total_annotations += ann_count
            print(f"{info['id']:<6} {info['name']:<35} {ann_count:<15}")

        print("-" * 60)
        print(f"{'합계':<42} {total_annotations:<15}")
        print("="*60 + "\n")

    def remove_classes(self, class_ids: List[Union[int, str]]) -> 'YOLOAnnotationEditor':
        """
        Remove specific classes

        Args:
            class_ids: List of class IDs or names to remove
        """
        names = self.config.get('names', {})

        # Convert names to IDs
        name_to_id = {name: cid for cid, name in names.items()}
        ids_to_remove = set()

        print("\n[변경 예정]")
        for item in class_ids:
            try:
                # Try as ID first
                class_id = int(item)
                if class_id in names:
                    ids_to_remove.add(class_id)
                    print(f"  - 제거: {names[class_id]} (ID: {class_id})")
                else:
                    print(f"[!] ID {class_id}를 찾을 수 없습니다")
            except ValueError:
                # Try as name
                if item in name_to_id:
                    class_id = name_to_id[item]
                    ids_to_remove.add(class_id)
                    print(f"  - 제거: {item} (ID: {class_id})")
                else:
                    print(f"[!] 이름 '{item}'을 찾을 수 없습니다")

        if not ids_to_remove:
            print("[!] 제거할 클래스를 찾을 수 없습니다")
            return self

        # Store statistics before change
        class_count_before = len(names)

        # Remove from data.yaml
        new_names = {cid: name for cid, name in names.items() if cid not in ids_to_remove}
        self.config['names'] = new_names

        # Remove annotations from label files
        total_removed = 0
        failed_files = []
        total_files = len(self.label_files)

        print(f"[*] {total_files}개 파일 처리 중...")
        for idx, label_file in enumerate(self.label_files, 1):
            if idx % 100 == 0 or idx == total_files:
                print(f"[*] 진행: {idx}/{total_files} ({idx*100//total_files}%)")
            try:
                with open(label_file, 'r', encoding='utf-8') as f:
                    lines = f.readlines()

                new_lines = []
                for line in lines:
                    if line.strip():
                        parts = line.split()
                        if len(parts) < 5:
                            print(f"[!] Warning: Skipping invalid line in {label_file}")
                            continue
                        try:
                            class_id = int(parts[0])
                            if class_id not in ids_to_remove:
                                new_lines.append(line)
                            else:
                                total_removed += 1
                        except ValueError:
                            print(f"[!] Warning: Skipping line with invalid class ID in {label_file}")
                            new_lines.append(line)

                if not self.dry_run:
                    with open(label_file, 'w', encoding='utf-8') as f:
                        f.writelines(new_lines)

            except PermissionError:
                print(f"[!] Error: Permission denied writing to {label_file}")
                failed_files.append(label_file)
            except Exception as e:
                print(f"[!] Error processing {label_file}: {e}")
                failed_files.append(label_file)

        if failed_files:
            print(f"[!] Warning: {len(failed_files)} files failed to process")

        # Print summary
        print(f"\n{'='*60}")
        print(f"작업 요약 - 클래스 제거")
        print(f"{'='*60}")
        print(f"제거된 클래스 수: {len(ids_to_remove)}")
        print(f"클래스 수 변화: {class_count_before} → {len(self.config['names'])}")
        print(f"제거된 어노테이션 수: {total_removed}")
        if failed_files:
            print(f"실패한 파일 수: {len(failed_files)}")
        print(f"{'='*60}\n")

        if self.dry_run:
            print(f"[DRY RUN] 위 변경사항이 적용될 예정입니다")
        else:
            print(f"[+] 작업이 완료되었습니다")
        return self

    def merge_classes(self, source_ids: List[Union[int, str]], target_name: str) -> 'YOLOAnnotationEditor':
        """
        Merge multiple classes into one

        Args:
            source_ids: List of source class IDs or names to merge
            target_name: Target class name after merge
        """
        names = self.config.get('names', {})
        name_to_id = {name: cid for cid, name in names.items()}

        # Find source class IDs
        source_class_ids = set()
        for item in source_ids:
            try:
                class_id = int(item)
                if class_id in names:
                    source_class_ids.add(class_id)
                    print(f"  - 통합 대상: {names[class_id]} (ID: {class_id})")
                else:
                    print(f"[!] ID {class_id}를 찾을 수 없습니다")
            except ValueError:
                if item in name_to_id:
                    class_id = name_to_id[item]
                    source_class_ids.add(class_id)
                    print(f"  - 통합 대상: {item} (ID: {class_id})")
                else:
                    print(f"[!] 이름 '{item}'을 찾을 수 없습니다")

        if not source_class_ids:
            print("[!] 통합할 클래스를 찾을 수 없습니다")
            return self

        # Use the smallest ID as the new class ID
        new_class_id = min(source_class_ids)

        # Store statistics before change
        class_count_before = len(names)

        # Update data.yaml
        new_names = {cid: name for cid, name in names.items() if cid not in source_class_ids}
        new_names[new_class_id] = target_name
        self.config['names'] = new_names

        # Update label files
        total_merged = 0
        failed_files = []
        total_files = len(self.label_files)

        print(f"[*] {total_files}개 파일 처리 중...")
        for idx, label_file in enumerate(self.label_files, 1):
            if idx % 100 == 0 or idx == total_files:
                print(f"[*] 진행: {idx}/{total_files} ({idx*100//total_files}%)")
            try:
                with open(label_file, 'r', encoding='utf-8') as f:
                    lines = f.readlines()

                new_lines = []
                for line in lines:
                    if line.strip():
                        parts = line.split()
                        if len(parts) < 5:
                            print(f"[!] Warning: Skipping invalid line in {label_file}")
                            continue
                        try:
                            class_id = int(parts[0])
                            if class_id in source_class_ids:
                                # Replace with new class ID
                                parts[0] = str(new_class_id)
                                new_lines.append(' '.join(parts) + '\n')
                                total_merged += 1
                            else:
                                new_lines.append(line)
                        except ValueError:
                            print(f"[!] Warning: Skipping line with invalid class ID in {label_file}")
                            new_lines.append(line)

                if not self.dry_run:
                    with open(label_file, 'w', encoding='utf-8') as f:
                        f.writelines(new_lines)

            except PermissionError:
                print(f"[!] Error: Permission denied writing to {label_file}")
                failed_files.append(label_file)
            except Exception as e:
                print(f"[!] Error processing {label_file}: {e}")
                failed_files.append(label_file)

        if failed_files:
            print(f"[!] Warning: {len(failed_files)} files failed to process")

        # Print summary
        print(f"\n{'='*60}")
        print(f"작업 요약 - 클래스 통합")
        print(f"{'='*60}")
        print(f"통합된 클래스 수: {len(source_class_ids)}")
        print(f"클래스 수 변화: {class_count_before} → {len(self.config['names'])}")
        print(f"새 클래스: {target_name} (ID: {new_class_id})")
        print(f"변경된 어노테이션 수: {total_merged}")
        if failed_files:
            print(f"실패한 파일 수: {len(failed_files)}")
        print(f"{'='*60}\n")

        if self.dry_run:
            print(f"[DRY RUN] 위 변경사항이 적용될 예정입니다")
        else:
            print(f"[+] 작업이 완료되었습니다")
        return self

    def rename_classes(self, mapping: Dict[Union[int, str], str]) -> 'YOLOAnnotationEditor':
        """
        Rename classes

        Args:
            mapping: Dictionary of {old_id_or_name: new_name}
        """
        names = self.config.get('names', {})
        name_to_id = {name: cid for cid, name in names.items()}

        renamed_count = 0
        for old_key, new_name in mapping.items():
            try:
                # Try as ID first
                class_id = int(old_key)
                if class_id in names:
                    old_name = names[class_id]
                    names[class_id] = new_name
                    print(f"  - 변경: '{old_name}' → '{new_name}' (ID: {class_id})")
                    renamed_count += 1
                else:
                    print(f"[!] ID {class_id}를 찾을 수 없습니다")
            except ValueError:
                # Try as name
                if old_key in name_to_id:
                    class_id = name_to_id[old_key]
                    names[class_id] = new_name
                    print(f"  - 변경: '{old_key}' → '{new_name}' (ID: {class_id})")
                    renamed_count += 1
                else:
                    print(f"[!] 이름 '{old_key}'를 찾을 수 없습니다")

        self.config['names'] = names
        if self.dry_run:
            print(f"[DRY RUN] {renamed_count}개 클래스 이름이 변경될 예정입니다")
        else:
            print(f"[+] {renamed_count}개 클래스 이름이 변경되었습니다")
        return self

    def reindex_classes(self, start_id: int = 0) -> 'YOLOAnnotationEditor':
        """
        Reindex class IDs sequentially

        Args:
            start_id: Starting ID (default: 0 for YOLO)
        """
        names = self.config.get('names', {})

        # Create ID mapping: old ID -> new ID
        id_mapping = {}
        new_id = start_id

        for old_id in sorted(names.keys()):
            id_mapping[old_id] = new_id
            print(f"  - ID 변경: {old_id} → {new_id} ({names[old_id]})")
            new_id += 1

        # Update data.yaml
        new_names = {id_mapping[old_id]: name for old_id, name in names.items()}
        self.config['names'] = new_names

        # Update label files
        total_updated = 0
        failed_files = []
        total_files = len(self.label_files)

        print(f"[*] {total_files}개 파일 처리 중...")
        for idx, label_file in enumerate(self.label_files, 1):
            if idx % 100 == 0 or idx == total_files:
                print(f"[*] 진행: {idx}/{total_files} ({idx*100//total_files}%)")
            try:
                with open(label_file, 'r', encoding='utf-8') as f:
                    lines = f.readlines()

                new_lines = []
                for line in lines:
                    if line.strip():
                        parts = line.split()
                        if len(parts) < 5:
                            print(f"[!] Warning: Skipping invalid line in {label_file}")
                            continue
                        try:
                            old_class_id = int(parts[0])
                            if old_class_id in id_mapping:
                                parts[0] = str(id_mapping[old_class_id])
                                new_lines.append(' '.join(parts) + '\n')
                                total_updated += 1
                            else:
                                new_lines.append(line)
                        except ValueError:
                            print(f"[!] Warning: Skipping line with invalid class ID in {label_file}")
                            new_lines.append(line)

                if not self.dry_run:
                    with open(label_file, 'w', encoding='utf-8') as f:
                        f.writelines(new_lines)

            except PermissionError:
                print(f"[!] Error: Permission denied writing to {label_file}")
                failed_files.append(label_file)
            except Exception as e:
                print(f"[!] Error processing {label_file}: {e}")
                failed_files.append(label_file)

        if failed_files:
            print(f"[!] Warning: {len(failed_files)} files failed to process")

        if self.dry_run:
            print(f"[DRY RUN] 클래스 ID가 재할당될 예정입니다 (시작: {start_id})")
            print(f"[DRY RUN] {total_updated}개의 어노테이션이 업데이트될 예정입니다")
        else:
            print(f"[+] 클래스 ID가 재할당되었습니다 (시작: {start_id})")
            print(f"[+] {total_updated}개의 어노테이션이 업데이트되었습니다")
        return self


def interactive_mode():
    """Interactive mode for user-friendly editing"""
    try:
        _interactive_mode_impl()
    except KeyboardInterrupt:
        print("\n\n[*] 사용자에 의해 중단되었습니다.")
        print("[*] 변경사항이 저장되지 않았습니다.")


def _interactive_mode_impl():
    """Internal implementation of interactive mode"""
    print("\n" + "="*60)
    print("  YOLO 어노테이션 편집기 - 대화형 모드")
    print("="*60)

    # 1. Select dataset
    print("\n[1/3] 데이터셋 선택")
    print("-" * 60)

    # Find YOLO datasets
    from glob import glob
    yaml_files = glob("**/data.yaml", recursive=True)

    dataset_path = None
    editor = None

    while not editor:
        if yaml_files:
            print("\n발견된 YOLO 데이터셋:")
            for idx, f in enumerate(yaml_files, 1):
                dataset_dir = str(Path(f).parent)
                print(f"  {idx}. {dataset_dir}")
            print(f"  0. 직접 경로 입력")
            print(f"  q. 종료")

            choice = input("\n데이터셋 번호 선택 (직접 입력은 0, 종료는 q): ").strip().lower()

            if choice == 'q':
                print("[*] 종료합니다.")
                return

            if choice == '0':
                dataset_path = input("데이터셋 디렉토리 경로 입력 (취소: q): ").strip()
                if dataset_path.lower() == 'q':
                    print("[*] 종료합니다.")
                    return
            else:
                try:
                    idx = int(choice)
                    if 1 <= idx <= len(yaml_files):
                        dataset_path = str(Path(yaml_files[idx - 1]).parent)
                    else:
                        print(f"[!] 1에서 {len(yaml_files)} 사이의 숫자를 입력해주세요.")
                        continue
                except ValueError:
                    print("[!] 잘못된 입력입니다. 숫자를 입력해주세요.")
                    continue
        else:
            dataset_path = input("데이터셋 디렉토리 경로 입력 (종료: q): ").strip()
            if dataset_path.lower() == 'q':
                print("[*] 종료합니다.")
                return

        if not dataset_path:
            print("[!] 경로를 입력해주세요.")
            continue

        # Try to load dataset
        try:
            print(f"\n[+] 로딩 중: {dataset_path}")
            editor = YOLOAnnotationEditor(dataset_path)
            editor.print_info()
        except FileNotFoundError as e:
            print(f"[!] 오류: {e}")
            print("[!] 다시 시도해주세요.")
            dataset_path = None
        except Exception as e:
            print(f"[!] 데이터셋 로딩 실패: {e}")
            print("[!] 다시 시도해주세요.")
            dataset_path = None

    # 2. Select operations
    print("\n[2/3] 작업 선택")
    print("-" * 60)
    operations = []

    while True:
        print("\n사용 가능한 작업:")
        print("  1. 클래스 제거")
        print("  2. 클래스 통합 (병합)")
        print("  3. 클래스 이름 변경")
        print("  4. 클래스 ID 재할당")
        print("  5. 현재 정보 표시")
        print("  0. 완료 (저장 단계로)")
        print("  h. 도움말")
        print("  q. 저장하지 않고 종료")

        choice = input("\n작업 선택 (0-5, h, q): ").strip().lower()

        if choice == 'q':
            confirm = input("저장하지 않고 종료하시겠습니까? (y/N): ").strip().lower()
            if confirm == 'y':
                print("[*] 종료합니다.")
                return
            continue

        if choice == '0':
            if not operations:
                print("[!] 선택된 작업이 없습니다.")
                confirm = input("그래도 종료하시겠습니까? (y/N): ").strip().lower()
                if confirm == 'y':
                    print("[*] 종료합니다.")
                    return
                continue
            break

        elif choice == '1':  # Remove
            print("\n현재 클래스:")
            for cid, name in sorted(editor.config['names'].items()):
                print(f"  {cid}: {name}")

            classes = input("\n제거할 클래스 입력 (ID 또는 이름, 쉼표로 구분, 취소: q): ").strip()
            if classes.lower() == 'q':
                print("[*] 작업을 취소했습니다.")
                continue
            if not classes:
                print("[!] 입력이 비어있습니다. 작업을 건너뜁니다.")
                continue

            class_list = [c.strip() for c in classes.split(',') if c.strip()]
            if not class_list:
                print("[!] 유효한 클래스가 없습니다.")
                continue

            confirm = input(f"\n정말로 {len(class_list)}개의 클래스를 제거하시겠습니까? (y/N): ").strip().lower()
            if confirm != 'y':
                print("[*] 작업을 취소했습니다.")
                continue

            print(f"\n[*] 클래스 제거 중: {class_list}")
            editor.remove_classes(class_list)
            operations.append(f"제거: {', '.join(class_list)}")

        elif choice == '2':  # Merge
            print("\n현재 클래스:")
            for cid, name in sorted(editor.config['names'].items()):
                print(f"  {cid}: {name}")

            source = input("\n통합할 클래스 입력 (ID 또는 이름, 쉼표로 구분, 취소: q): ").strip()
            if source.lower() == 'q':
                print("[*] 작업을 취소했습니다.")
                continue
            if not source:
                print("[!] 입력이 비어있습니다. 작업을 건너뜁니다.")
                continue

            source_list = [c.strip() for c in source.split(',') if c.strip()]
            if len(source_list) < 2:
                print("[!] 최소 2개 이상의 클래스를 입력해야 합니다.")
                continue

            target = input("통합 후 클래스 이름 입력 (취소: q): ").strip()
            if target.lower() == 'q':
                print("[*] 작업을 취소했습니다.")
                continue
            if not target:
                print("[!] 클래스 이름을 입력해주세요.")
                continue

            confirm = input(f"\n{len(source_list)}개의 클래스를 '{target}'로 통합하시겠습니까? (y/N): ").strip().lower()
            if confirm != 'y':
                print("[*] 작업을 취소했습니다.")
                continue

            print(f"\n[*] '{target}'로 통합 중: {source_list}")
            editor.merge_classes(source_list, target)
            operations.append(f"통합: {', '.join(source_list)} → {target}")

        elif choice == '3':  # Rename
            print("\n현재 클래스:")
            for cid, name in sorted(editor.config['names'].items()):
                print(f"  {cid}: {name}")

            print("\n매핑 입력 (형식: 기존:새이름,기존:새이름 - 쉼표로 구분)")
            print("예시: 0:vehicle,1:person 또는 truck:vehicle,person:human")
            mapping_input = input("\n입력 (취소: q): ").strip()

            if mapping_input.lower() == 'q':
                print("[*] 작업을 취소했습니다.")
                continue
            if not mapping_input:
                print("[!] 입력이 비어있습니다. 작업을 건너뜁니다.")
                continue

            mapping = {}
            invalid_pairs = []
            for pair in mapping_input.split(','):
                pair = pair.strip()
                if ':' in pair:
                    parts = pair.split(':', 1)
                    old, new = parts[0].strip(), parts[1].strip()
                    if old and new:
                        mapping[old] = new
                    else:
                        invalid_pairs.append(pair)
                else:
                    invalid_pairs.append(pair)

            if invalid_pairs:
                print(f"[!] 잘못된 형식의 매핑: {', '.join(invalid_pairs)}")

            if not mapping:
                print("[!] 유효한 매핑이 없습니다.")
                continue

            print(f"\n변경 예정:")
            for old, new in mapping.items():
                print(f"  {old} → {new}")

            confirm = input(f"\n{len(mapping)}개의 클래스 이름을 변경하시겠습니까? (y/N): ").strip().lower()
            if confirm != 'y':
                print("[*] 작업을 취소했습니다.")
                continue

            print(f"\n[*] 클래스 이름 변경 중")
            editor.rename_classes(mapping)
            operations.append(f"이름 변경: {len(mapping)}개 클래스")

        elif choice == '4':  # Reindex
            while True:
                start_id_input = input("\n시작 ID 입력 (기본값: 0, 취소: q): ").strip().lower()
                if start_id_input == 'q':
                    print("[*] 작업을 취소했습니다.")
                    start_id = None
                    break
                if not start_id_input:
                    start_id = 0
                    break
                try:
                    start_id = int(start_id_input)
                    if start_id < 0:
                        print("[!] 0 이상의 숫자를 입력해주세요.")
                        continue
                    break
                except ValueError:
                    print("[!] 유효한 숫자를 입력해주세요.")

            if start_id is None:
                continue

            confirm = input(f"\n클래스 ID를 {start_id}부터 재할당하시겠습니까? (y/N): ").strip().lower()
            if confirm != 'y':
                print("[*] 작업을 취소했습니다.")
                continue

            print(f"\n[*] 클래스 ID 재할당 중 (시작: {start_id})")
            editor.reindex_classes(start_id=start_id)
            operations.append(f"ID 재할당: {start_id}부터")

        elif choice == '5':  # Show info
            editor.print_info()

        elif choice == 'h':  # Help
            print("\n" + "="*60)
            print("  도움말")
            print("="*60)
            print("\n1. 클래스 제거")
            print("   - 지정한 클래스와 해당 어노테이션을 완전히 제거합니다")
            print("   - 예: '0,1' 또는 'person,car'")
            print("\n2. 클래스 통합 (병합)")
            print("   - 여러 클래스를 하나로 합칩니다")
            print("   - 예: car,truck,bus → vehicle")
            print("\n3. 클래스 이름 변경")
            print("   - 클래스 이름만 변경합니다 (ID는 유지)")
            print("   - 예: 0:vehicle,1:person")
            print("\n4. 클래스 ID 재할당")
            print("   - 모든 클래스 ID를 순차적으로 재할당합니다")
            print("   - 예: 시작 ID 0 → 0,1,2,3...")
            print("\n5. 현재 정보 표시")
            print("   - 데이터셋의 현재 상태를 확인합니다")
            print("\n팁:")
            print("   - 모든 작업은 적용 전 확인 메시지가 표시됩니다")
            print("   - 'q'를 입력하여 언제든지 작업을 취소할 수 있습니다")
            print("   - 저장 전 백업이 자동으로 생성됩니다")
            print("="*60)

        else:
            print("[!] 잘못된 선택입니다")

    # 3. Save
    print("\n[3/3] 결과 저장")
    print("-" * 60)
    print("\n수행된 작업:")
    for idx, op in enumerate(operations, 1):
        print(f"  {idx}. {op}")

    # Ask for labels backup
    print("\n[!] 주의: labels 폴더의 모든 .txt 파일이 수정됩니다.")
    create_labels_backup = input("labels 폴더 백업 생성? (Y/n): ").strip().lower()
    if create_labels_backup != 'n':
        editor.backup_labels()

    create_yaml_backup = input("data.yaml 백업 생성? (Y/n): ").strip().lower()
    yaml_backup = create_yaml_backup != 'n'

    print(f"\n[*] 저장 중...")
    editor.save_yaml(backup=yaml_backup)

    print("\n" + "="*60)
    print("  작업이 성공적으로 완료되었습니다!")
    print("="*60)
    print(f"\n데이터셋: {dataset_path}")
    print(f"설정 파일: {editor.yaml_path}")
    print(f"라벨 파일: {len(editor.label_files)}개 업데이트됨")
    print()


def main():
    parser = argparse.ArgumentParser(
        description='YOLO Annotation Editing Utility',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__
    )

    subparsers = parser.add_subparsers(dest='command', help='Available commands')

    # info command
    info_parser = subparsers.add_parser('info', help='Show annotation info')
    info_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')

    # remove command
    remove_parser = subparsers.add_parser('remove', help='Remove classes')
    remove_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')
    remove_parser.add_argument('-c', '--classes', required=True, help='Classes to remove (comma-separated IDs or names)')
    remove_parser.add_argument('--no-backup', action='store_true', help='Do not create backup for labels and yaml')
    remove_parser.add_argument('--no-labels-backup', action='store_true', help='Do not create backup for labels folder')

    # merge command
    merge_parser = subparsers.add_parser('merge', help='Merge classes')
    merge_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')
    merge_parser.add_argument('-s', '--source', required=True, help='Source classes to merge (comma-separated)')
    merge_parser.add_argument('-t', '--target', required=True, help='Target class name after merge')
    merge_parser.add_argument('--no-backup', action='store_true', help='Do not create backup for labels and yaml')
    merge_parser.add_argument('--no-labels-backup', action='store_true', help='Do not create backup for labels folder')

    # rename command
    rename_parser = subparsers.add_parser('rename', help='Rename classes')
    rename_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')
    rename_parser.add_argument('-m', '--mapping', required=True,
                              help='Name mapping (old:new,old2:new2 format)')
    rename_parser.add_argument('--no-backup', action='store_true', help='Do not create backup for labels and yaml')
    rename_parser.add_argument('--no-labels-backup', action='store_true', help='Do not create backup for labels folder')

    # reindex command
    reindex_parser = subparsers.add_parser('reindex', help='Reindex class IDs')
    reindex_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')
    reindex_parser.add_argument('--start', type=int, default=0, help='Starting ID (default: 0)')
    reindex_parser.add_argument('--no-backup', action='store_true', help='Do not create backup for labels and yaml')
    reindex_parser.add_argument('--no-labels-backup', action='store_true', help='Do not create backup for labels folder')

    # If no arguments provided, enter interactive mode
    if len(sys.argv) == 1:
        interactive_mode()
        return

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return

    # Execute command
    if args.command == 'info':
        editor = YOLOAnnotationEditor(args.dataset)
        editor.print_info()

    elif args.command == 'remove':
        classes = [c.strip() for c in args.classes.split(',')]
        editor = YOLOAnnotationEditor(args.dataset)

        # Backup labels if requested
        if not args.no_backup and not args.no_labels_backup:
            editor.backup_labels()

        print(f"\n[*] 클래스 제거 중...")
        editor.remove_classes(classes)
        editor.save_yaml(backup=not args.no_backup)

    elif args.command == 'merge':
        source_classes = [c.strip() for c in args.source.split(',')]
        editor = YOLOAnnotationEditor(args.dataset)

        # Backup labels if requested
        if not args.no_backup and not args.no_labels_backup:
            editor.backup_labels()

        print(f"\n[*] 클래스 통합 중...")
        editor.merge_classes(source_classes, args.target)
        editor.save_yaml(backup=not args.no_backup)

    elif args.command == 'rename':
        mapping = {}
        for pair in args.mapping.split(','):
            old, new = pair.split(':')
            mapping[old.strip()] = new.strip()
        editor = YOLOAnnotationEditor(args.dataset)

        # Backup labels if requested
        if not args.no_backup and not args.no_labels_backup:
            editor.backup_labels()

        print(f"\n[*] 클래스 이름 변경 중...")
        editor.rename_classes(mapping)
        editor.save_yaml(backup=not args.no_backup)

    elif args.command == 'reindex':
        editor = YOLOAnnotationEditor(args.dataset)

        # Backup labels if requested
        if not args.no_backup and not args.no_labels_backup:
            editor.backup_labels()

        print(f"\n[*] 클래스 ID 재할당 중...")
        editor.reindex_classes(start_id=args.start)
        editor.save_yaml(backup=not args.no_backup)


if __name__ == '__main__':
    main()