import argparse
from enum import Enum


class Model(Enum):
    AST = "AST"
    Classical = "Classical"
    HuBERT = "HuBERT"
    RawNet2 = "RawNet2"
    VIT = "VIT"
    Wav2Vec2 = "Wav2Vec2"
    WavLM = "WavLM"


class Dataset(Enum):
    ASVspoof2019 = "ASVspoof2019"
    ASVspoof5 = "ASVspoof5"
    VoxCelebSpoof = "VoxCelebSpoof"


class Visualisation(Enum):
    ConstantQ = "ConstantQ"
    MelSpectrogram = "MelSpectrogram"
    MFCC = "MFCC"


class EnumAction(argparse.Action):
    def __init__(self, **kwargs):
        enum_type = kwargs.pop("type", None)
        if enum_type is None:
            raise ValueError("type must be assigned an Enum when using EnumAction")
        if not issubclass(enum_type, Enum):
            raise TypeError("type must be an Enum when using EnumAction")
        kwargs.setdefault("choices", tuple(e.name for e in enum_type))

        super(EnumAction, self).__init__(**kwargs)

        self._enum = enum_type

    def __call__(self, parser, namespace, values, option_string=None):
        try:
            value = self._enum[values]
        except KeyError:
            choices = ", ".join(e.name for e in self._enum)
            parser.error(f"invalid choice '{values}' — choose from: {choices}")
        setattr(namespace, self.dest, value)
