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

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

    # Remove classes
    python mmdet_annotation_modify.py remove --input annotations.json --output new.json --classes "person,car"

    # Merge classes
    python mmdet_annotation_modify.py merge --input annotations.json --output new.json --source "car,truck,bus" --target "vehicle"

    # Rename classes
    python mmdet_annotation_modify.py rename --input annotations.json --output new.json --mapping "old_name:new_name,another_old:another_new"

    # Show info
    python mmdet_annotation_modify.py info --input annotations.json
"""

import json
import argparse
from pathlib import Path
from typing import Dict, List, Set, Union
import shutil
from datetime import datetime
import sys
import os


class COCOAnnotationEditor:
    """Class for editing COCO format annotations"""

    def __init__(self, annotation_path: str, dry_run: bool = False):
        """
        Args:
            annotation_path: Path to COCO format annotation JSON file
            dry_run: If True, only simulate changes without writing files
        """
        self.annotation_path = Path(annotation_path)
        self.dry_run = dry_run

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

        # Validate file exists
        if not self.annotation_path.exists():
            raise FileNotFoundError(f"Annotation file not found: {annotation_path}")

        if not self.annotation_path.is_file():
            raise ValueError(f"Path is not a file: {annotation_path}")

        self.data = self._load_annotation()

    def _load_annotation(self) -> dict:
        """Load annotation file"""
        try:
            with open(self.annotation_path, 'r', encoding='utf-8') as f:
                data = json.load(f)

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

            required_fields = ['images', 'annotations', 'categories']
            for field in required_fields:
                if field not in data:
                    raise ValueError(f"Missing required field: '{field}'")

            if not isinstance(data['categories'], list):
                raise ValueError("'categories' field must be a list")

            if not isinstance(data['annotations'], list):
                raise ValueError("'annotations' field must be a list")

            if not isinstance(data['images'], list):
                raise ValueError("'images' field must be a list")

            return data

        except json.JSONDecodeError as e:
            raise ValueError(f"Failed to parse JSON file: {e}")
        except UnicodeDecodeError:
            raise ValueError(f"Cannot decode file (invalid encoding)")
        except Exception as e:
            raise RuntimeError(f"Error loading annotation file: {e}")

    def save(self, output_path: str, backup: bool = True):
        """
        Save modified annotation

        Args:
            output_path: Output file path
            backup: Whether to backup original file
        """
        output_path = Path(output_path)

        if self.dry_run:
            print(f"[DRY RUN] 파일이 저장될 예정입니다: {output_path}")
            if backup and output_path.exists():
                print(f"[DRY RUN] 백업이 생성될 예정입니다")
            return

        try:
            # Create backup
            if backup and output_path.exists():
                backup_path = output_path.parent / f"{output_path.stem}_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}{output_path.suffix}"
                shutil.copy2(output_path, backup_path)
                print(f"[+] 백업 생성됨: {backup_path}")

            # Ensure parent directory exists
            output_path.parent.mkdir(parents=True, exist_ok=True)

            # Save
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(self.data, f, ensure_ascii=False, indent=2)
            print(f"[+] 저장 완료: {output_path}")

        except PermissionError:
            raise PermissionError(f"Permission denied writing to {output_path}")
        except Exception as e:
            raise RuntimeError(f"Failed to save annotation file: {e}")

    def get_category_info(self) -> Dict:
        """Get category information"""
        categories = self.data.get('categories', [])
        annotations = self.data.get('annotations', [])

        # Count annotations per category
        category_counts = {}
        for idx, ann in enumerate(annotations):
            if 'category_id' not in ann:
                print(f"[!] Warning: Annotation {idx} missing 'category_id' field")
                continue
            try:
                cat_id = ann['category_id']
                category_counts[cat_id] = category_counts.get(cat_id, 0) + 1
            except Exception as e:
                print(f"[!] Warning: Error processing annotation {idx}: {e}")

        info = {}
        for cat in categories:
            if 'id' not in cat or 'name' not in cat:
                print(f"[!] Warning: Category missing required fields: {cat}")
                continue
            try:
                cat_id = cat['id']
                info[cat_id] = {
                    'id': cat_id,
                    'name': cat['name'],
                    'supercategory': cat.get('supercategory', ''),
                    'annotation_count': category_counts.get(cat_id, 0)
                }
            except Exception as e:
                print(f"[!] Warning: Error processing category: {e}")

        return info

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

        print(f"\n전체 통계:")
        print(f"  - 이미지 수: {len(self.data.get('images', []))}")
        print(f"  - 어노테이션 수: {len(self.data.get('annotations', []))}")
        print(f"  - 클래스 수: {len(self.data.get('categories', []))}")

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

        category_info = self.get_category_info()
        total_annotations = 0
        for cat_id, info in sorted(category_info.items()):
            ann_count = info['annotation_count']
            total_annotations += ann_count
            print(f"{info['id']:<6} {info['name']:<30} {ann_count:<15}")

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

    def remove_categories(self, category_names: List[str]) -> 'COCOAnnotationEditor':
        """
        Remove specific categories

        Args:
            category_names: List of category names or IDs to remove
        """
        # Build ID to name mapping
        id_to_cat = {cat['id']: cat for cat in self.data.get('categories', [])}
        name_to_cat = {cat['name']: cat for cat in self.data.get('categories', [])}

        # Find category IDs to remove
        categories_to_remove = set()

        for item in category_names:
            # Try to parse as ID first
            try:
                cat_id = int(item)
                if cat_id in id_to_cat:
                    cat = id_to_cat[cat_id]
                    categories_to_remove.add(cat['id'])
                    print(f"  - 제거: {cat['name']} (ID: {cat['id']})")
                else:
                    print(f"[!] ID {cat_id}를 찾을 수 없습니다")
            except ValueError:
                # Not a number, treat as name
                if item in name_to_cat:
                    cat = name_to_cat[item]
                    categories_to_remove.add(cat['id'])
                    print(f"  - 제거: {cat['name']} (ID: {cat['id']})")
                else:
                    print(f"[!] 이름 '{item}'을 찾을 수 없습니다")

        # Update categories
        remaining_categories = [
            cat for cat in self.data.get('categories', [])
            if cat['id'] not in categories_to_remove
        ]

        # Update categories
        self.data['categories'] = remaining_categories

        # Remove annotations of removed categories
        original_count = len(self.data.get('annotations', []))
        print(f"[*] {original_count}개 어노테이션 검사 중...")

        new_annotations = []
        for idx, ann in enumerate(self.data.get('annotations', []), 1):
            if idx % 1000 == 0:
                print(f"[*] 진행: {idx}/{original_count} ({idx*100//original_count}%)")
            if ann['category_id'] not in categories_to_remove:
                new_annotations.append(ann)

        self.data['annotations'] = new_annotations
        removed_count = original_count - len(new_annotations)

        if self.dry_run:
            print(f"[DRY RUN] {removed_count}개의 어노테이션이 제거될 예정입니다")
        else:
            print(f"[+] {removed_count}개의 어노테이션이 제거되었습니다")
        return self


    def merge_categories(self, source_names: List[str], target_name: str,
                        target_supercategory: str = '') -> 'COCOAnnotationEditor':
        """
        Merge multiple categories into one

        Args:
            source_names: List of source category names or IDs to merge
            target_name: Target category name after merge
            target_supercategory: Target supercategory (optional)
        """
        # Find source category IDs
        source_ids = set()
        source_categories = []

        # Build ID to name mapping
        id_to_cat = {cat['id']: cat for cat in self.data.get('categories', [])}
        name_to_cat = {cat['name']: cat for cat in self.data.get('categories', [])}

        for item in source_names:
            # Try to parse as ID first
            try:
                cat_id = int(item)
                if cat_id in id_to_cat:
                    cat = id_to_cat[cat_id]
                    source_ids.add(cat['id'])
                    source_categories.append(cat)
                    print(f"  - 통합 대상: {cat['name']} (ID: {cat['id']})")
                else:
                    print(f"[!] ID {cat_id}를 찾을 수 없습니다")
            except ValueError:
                # Not a number, treat as name
                if item in name_to_cat:
                    cat = name_to_cat[item]
                    source_ids.add(cat['id'])
                    source_categories.append(cat)
                    print(f"  - 통합 대상: {cat['name']} (ID: {cat['id']})")
                else:
                    print(f"[!] 이름 '{item}'을 찾을 수 없습니다")

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

        # New category ID (use first source ID)
        new_category_id = min(source_ids)

        # Create new category
        new_category = {
            'id': new_category_id,
            'name': target_name,
            'supercategory': ''
        }

        # Remove source categories and add new category
        self.data['categories'] = [
            cat for cat in self.data.get('categories', [])
            if cat['id'] not in source_ids
        ]
        self.data['categories'].append(new_category)

        # Update annotations (change all source IDs to new ID)
        total_anns = len(self.data.get('annotations', []))
        print(f"[*] {total_anns}개 어노테이션 업데이트 중...")

        for idx, ann in enumerate(self.data.get('annotations', []), 1):
            if idx % 1000 == 0:
                print(f"[*] 진행: {idx}/{total_anns} ({idx*100//total_anns}%)")
            if ann['category_id'] in source_ids:
                ann['category_id'] = new_category_id

        if self.dry_run:
            print(f"[DRY RUN] '{target_name}' (ID: {new_category_id})로 통합될 예정입니다")
        else:
            print(f"[+] '{target_name}' (ID: {new_category_id})로 통합되었습니다")
        return self

    def rename_categories(self, mapping: Dict[str, str]) -> 'COCOAnnotationEditor':
        """
        Rename categories

        Args:
            mapping: Dictionary of {old_name_or_id: new_name}
        """
        # Build ID to cat mapping
        id_to_cat = {cat['id']: cat for cat in self.data.get('categories', [])}
        name_to_cat = {cat['name']: cat for cat in self.data.get('categories', [])}

        renamed_count = 0
        for old_key, new_name in mapping.items():
            # Try to parse as ID first
            try:
                cat_id = int(old_key)
                if cat_id in id_to_cat:
                    cat = id_to_cat[cat_id]
                    old_name = cat['name']
                    cat['name'] = new_name
                    print(f"  - 변경: '{old_name}' → '{new_name}' (ID: {cat['id']})")
                    renamed_count += 1
                else:
                    print(f"[!] ID {cat_id}를 찾을 수 없습니다")
            except ValueError:
                # Not a number, treat as name
                if old_key in name_to_cat:
                    cat = name_to_cat[old_key]
                    cat['name'] = new_name
                    print(f"  - 변경: '{old_key}' → '{new_name}' (ID: {cat['id']})")
                    renamed_count += 1
                else:
                    print(f"[!] 이름 '{old_key}'를 찾을 수 없습니다")

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

    def reindex_categories(self, start_id: int = 1) -> 'COCOAnnotationEditor':
        """
        Reindex category IDs sequentially

        Args:
            start_id: Starting ID (default: 1)
        """
        # Create mapping: old ID -> new ID
        id_mapping = {}
        new_id = start_id

        for cat in sorted(self.data.get('categories', []), key=lambda x: x['id']):
            old_id = cat['id']
            id_mapping[old_id] = new_id
            print(f"  - ID 변경: {old_id} → {new_id} ({cat['name']})")
            new_id += 1

        # Update category IDs
        for cat in self.data.get('categories', []):
            cat['id'] = id_mapping[cat['id']]

        # Update annotation category IDs
        total_anns = len(self.data.get('annotations', []))
        print(f"[*] {total_anns}개 어노테이션 업데이트 중...")

        for idx, ann in enumerate(self.data.get('annotations', []), 1):
            if idx % 1000 == 0:
                print(f"[*] 진행: {idx}/{total_anns} ({idx*100//total_anns}%)")
            ann['category_id'] = id_mapping[ann['category_id']]

        if self.dry_run:
            print(f"[DRY RUN] 클래스 ID가 재할당될 예정입니다 (시작: {start_id})")
        else:
            print(f"[+] 클래스 ID가 재할당되었습니다 (시작: {start_id})")
        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("  COCO 어노테이션 편집기 - 대화형 모드")
    print("="*60)

    # 1. Select annotation file
    print("\n[1/3] 어노테이션 파일 선택")
    print("-" * 60)

    # Find annotation files in common locations
    common_paths = [
        "dataset/*/annotations/*.json",
        "data/*/annotations/*.json",
        "annotations/*.json",
        "*.json"
    ]

    from glob import glob
    found_files = []
    for pattern in common_paths:
        found_files.extend(glob(pattern, recursive=True))

    input_file = None
    editor = None

    while not editor:
        if found_files:
            print("\n발견된 어노테이션 파일:")
            for idx, f in enumerate(found_files[:10], 1):
                print(f"  {idx}. {f}")
            if len(found_files) > 10:
                print(f"  ... 외 {len(found_files) - 10}개 더")
            print(f"  0. 직접 경로 입력")
            print(f"  q. 종료")

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

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

            if choice == '0':
                input_file = input("어노테이션 파일 경로 입력 (취소: q): ").strip()
                if input_file.lower() == 'q':
                    print("[*] 종료합니다.")
                    return
            else:
                try:
                    idx = int(choice)
                    if 1 <= idx <= len(found_files):
                        input_file = found_files[idx - 1]
                    else:
                        print(f"[!] 1에서 {len(found_files)} 사이의 숫자를 입력해주세요.")
                        continue
                except ValueError:
                    print("[!] 잘못된 입력입니다. 숫자를 입력해주세요.")
                    continue
        else:
            input_file = input("어노테이션 파일 경로 입력 (종료: q): ").strip()
            if input_file.lower() == 'q':
                print("[*] 종료합니다.")
                return

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

        # Try to load annotation
        try:
            print(f"\n[+] 로딩 중: {input_file}")
            editor = COCOAnnotationEditor(input_file)
            editor.print_info()
        except FileNotFoundError as e:
            print(f"[!] 오류: {e}")
            print("[!] 다시 시도해주세요.")
            input_file = None
        except Exception as e:
            print(f"[!] 어노테이션 파일 로딩 실패: {e}")
            print("[!] 다시 시도해주세요.")
            input_file = None

    # Store original categories for reference
    original_categories = {cat['id']: cat['name'] for cat in editor.data.get('categories', [])}

    # 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 cat in editor.data.get('categories', []):
                print(f"  - {cat['name']} (ID: {cat['id']})")

            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_categories(class_list)
            operations.append(f"제거: {', '.join(class_list)}")

        elif choice == '2':  # Merge
            print("\n현재 클래스:")
            for cat in editor.data.get('categories', []):
                print(f"  - {cat['name']} (ID: {cat['id']})")

            source = input("\n통합할 클래스 입력 (이름 또는 ID, 쉼표로 구분): ").strip()
            if source:
                target = input("통합 후 클래스 이름 입력: ").strip()

                source_list = [c.strip() for c in source.split(',')]
                print(f"\n[*] '{target}'로 통합 중: {source_list}")
                editor.merge_categories(source_list, target, '')
                operations.append(f"통합: {', '.join(source_list)} → {target}")

        elif choice == '3':  # Rename
            print("\n현재 클래스:")
            for cat in editor.data.get('categories', []):
                print(f"  - {cat['name']} (ID: {cat['id']})")

            print("\n매핑 입력 (형식: 기존:새이름,기존:새이름 - 쉼표로 구분)")
            print("예시: tt:transformer,ss:substation 또는 1:transformer,2:substation")
            mapping_input = input("\n입력: ").strip()

            if mapping_input:
                mapping = {}
                for pair in mapping_input.split(','):
                    if ':' in pair:
                        old, new = pair.split(':', 1)
                        mapping[old.strip()] = new.strip()

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

        elif choice == '4':  # Reindex
            while True:
                start_id = input("\n시작 ID 입력 (기본값: 1): ").strip()
                if not start_id:
                    start_id = 1
                    break
                try:
                    start_id = int(start_id)
                    break
                except ValueError:
                    print("[!] 숫자를 입력해주세요.")

            print(f"\n[*] 클래스 ID 재할당 중 (시작: {start_id})")
            editor.reindex_categories(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("   - 예: 'person,car' 또는 '1,2'")
            print("\n2. 클래스 통합 (병합)")
            print("   - 여러 클래스를 하나로 합칩니다")
            print("   - 예: car,truck,bus → vehicle")
            print("\n3. 클래스 이름 변경")
            print("   - 클래스 이름만 변경합니다 (ID는 유지)")
            print("   - 예: transformer:tt,substation:ss")
            print("\n4. 클래스 ID 재할당")
            print("   - 모든 클래스 ID를 순차적으로 재할당합니다")
            print("   - COCO 형식은 보통 1부터 시작합니다")
            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}")

    print(f"\n입력 파일: {input_file}")
    default_output = str(Path(input_file).parent / f"{Path(input_file).stem}_edited{Path(input_file).suffix}")
    output_file = input(f"\n출력 파일 경로 입력 (기본값: {default_output}): ").strip()
    if not output_file:
        output_file = default_output

    create_backup = input("출력 파일이 존재하면 백업 생성? (Y/n): ").strip().lower()
    backup = create_backup != 'n'

    print(f"\n[*] 저장 중: {output_file}...")
    editor.save(output_file, backup=backup)

    print("\n" + "="*60)
    print("  작업이 성공적으로 완료되었습니다!")
    print("="*60)
    print(f"\n입력:  {input_file}")
    print(f"출력: {output_file}")
    if backup and Path(output_file).exists():
        print(f"(파일이 존재했다면 백업이 생성되었습니다)")
    print()


def main():
    parser = argparse.ArgumentParser(
        description='COCO 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('-i', '--input', required=True, help='Input annotation file')

    # remove command
    remove_parser = subparsers.add_parser('remove', help='Remove categories')
    remove_parser.add_argument('-i', '--input', required=True, help='Input annotation file')
    remove_parser.add_argument('-o', '--output', required=True, help='Output annotation file')
    remove_parser.add_argument('-c', '--classes', required=True, help='Classes to remove (comma-separated)')
    remove_parser.add_argument('--no-backup', action='store_true', help='Do not create backup')

    # merge command
    merge_parser = subparsers.add_parser('merge', help='Merge categories')
    merge_parser.add_argument('-i', '--input', required=True, help='Input annotation file')
    merge_parser.add_argument('-o', '--output', required=True, help='Output annotation file')
    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('--supercategory', default='', help='Supercategory for merged class')
    merge_parser.add_argument('--no-backup', action='store_true', help='Do not create backup')

    # rename command
    rename_parser = subparsers.add_parser('rename', help='Rename categories')
    rename_parser.add_argument('-i', '--input', required=True, help='Input annotation file')
    rename_parser.add_argument('-o', '--output', required=True, help='Output annotation file')
    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')

    # reindex command
    reindex_parser = subparsers.add_parser('reindex', help='Reindex category IDs')
    reindex_parser.add_argument('-i', '--input', required=True, help='Input annotation file')
    reindex_parser.add_argument('-o', '--output', required=True, help='Output annotation file')
    reindex_parser.add_argument('--start', type=int, default=1, help='Starting ID (default: 1)')
    reindex_parser.add_argument('--no-backup', action='store_true', help='Do not create backup')

    # 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 = COCOAnnotationEditor(args.input)
        editor.print_info()

    elif args.command == 'remove':
        classes = [c.strip() for c in args.classes.split(',')]
        print(f"\n[*] 클래스 제거 중...")
        editor = COCOAnnotationEditor(args.input)
        editor.remove_categories(classes)
        editor.save(args.output, backup=not args.no_backup)

    elif args.command == 'merge':
        source_classes = [c.strip() for c in args.source.split(',')]
        print(f"\n[*] 클래스 통합 중...")
        editor = COCOAnnotationEditor(args.input)
        editor.merge_categories(source_classes, args.target, '')
        editor.save(args.output, 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()
        print(f"\n[*] 클래스 이름 변경 중...")
        editor = COCOAnnotationEditor(args.input)
        editor.rename_categories(mapping)
        editor.save(args.output, backup=not args.no_backup)

    elif args.command == 'reindex':
        print(f"\n[*] 클래스 ID 재할당 중...")
        editor = COCOAnnotationEditor(args.input)
        editor.reindex_categories(start_id=args.start)
        editor.save(args.output, backup=not args.no_backup)


if __name__ == '__main__':
    main()