import os
import sys
from pathlib import Path

def filter_illustrations_only(target_path: Path, inplace: bool = False) -> None:
    """
    Filtre les annotations YOLO pour ne garder que la classe 'Illustration' (0).
    Supprime toutes les autres classes (Initial, Ornament, Stamp, Table).
    """
    # Mapping : seule la classe 0 (Illustration) est conservée
    label_mapping = {
        0: 0,     # Illustration -> Gardée
        1: None,  # Initial -> Supprimée
        2: None,  # Ornament -> Supprimée
        3: None,  # Stamp -> Supprimée
        4: None,  # Table -> Supprimée
    }

    # Détection automatique du dossier 'labels'
    if target_path.name == "labels":
        input_dir = target_path
        output_dir = target_path if inplace else target_path.parent / "labels_output"
    else:
        input_dir = target_path / "labels"
        output_dir = input_dir if inplace else target_path / "labels_output"

    if not input_dir.exists():
        print(f"Erreur : Le dossier '{input_dir}' n'existe pas.")
        return

    if not inplace:
        os.makedirs(output_dir, exist_ok=True)

    total_files = 0
    total_kept = 0
    total_deleted = 0
    empty_files_count = 0

    for filename in os.listdir(input_dir):
        if not filename.endswith(".txt"):
            continue

        total_files += 1
        input_file = input_dir / filename
        output_file = output_dir / filename

        with open(input_file, "r", encoding="utf-8") as f:
            lines = f.readlines()

        new_lines = []
        for line in lines:
            parts = line.strip().split()
            if not parts:
                continue

            try:
                original_label = int(parts[0])
            except ValueError:
                continue  # Ligne mal formée, ignorée

            new_label = label_mapping.get(original_label)

            if new_label is None:
                total_deleted += 1
                continue

            parts[0] = str(new_label)
            new_lines.append(" ".join(parts))
            total_kept += 1

        # Écriture du fichier (vide si aucune illustration)
        with open(output_file, "w", encoding="utf-8") as f:
            if new_lines:
                f.write("\n".join(new_lines) + "\n")
            else:
                empty_files_count += 1
                f.write("")

    print("--- Résumé du traitement ---")
    print(f"Fichiers traités               : {total_files}")
    print(f"Annotations conservées (Illu) : {total_kept}")
    print(f"Annotations supprimées         : {total_deleted}")
    print(f"Fichiers .txt résultants vides : {empty_files_count}")
    print(f"Dossier de sortie              : {output_dir.resolve()}")
    print("Traitement terminé avec succès !")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: uv run filter_illustration_only.py <chemin_du_dossier>")
        sys.exit(1)

    dataset_path = Path(sys.argv[1])
    
    # mets inplace=True si tu veux écraser les fichiers d'origine dans 'labels/'
    filter_illustrations_only(dataset_path, inplace=False)