from pathlib import Path
from typing import List


DEFAULT_LABELS = [
    "cat",
    "dog",
    "rabbit",
]


PROJECT_DIR = Path(__file__).resolve().parent


def load_labels(dataset_train_dir: str = "data/custom_dataset/train", labels_file: str = "labels.txt") -> List[str]:
    """Load class labels from labels.txt or infer them from train subfolders."""
    labels_path = Path(labels_file)
    if not labels_path.is_absolute():
        labels_path = PROJECT_DIR / labels_path

    if labels_path.exists():
        labels = [line.strip() for line in labels_path.read_text(encoding="utf-8").splitlines() if line.strip()]
        if labels:
            return labels

    train_path = Path(dataset_train_dir)
    if not train_path.is_absolute():
        train_path = PROJECT_DIR / train_path

    if train_path.exists() and train_path.is_dir():
        inferred = sorted([p.name for p in train_path.iterdir() if p.is_dir()])
        if inferred:
            return inferred

    return DEFAULT_LABELS