"""Modifier definitions and parsing for voice generation."""

import copy
import enum
from dataclasses import dataclass
from typing import Dict, List, Protocol, Type, TypeAlias, Union

import librosa

from .voice_audio import Audio


class ModifierFormatNotCorrect(Exception):
    """Modifiers not provided in the correct format."""


class AcceleratorDirectionInvalid(Exception):
    """Invalid accelerator direction."""


class ModifierArgumentsInvalid(Exception):
    """Invalid modifier arguments and/or format."""


class ModifierClass(Protocol):
    IDENTIFIER: str

    @classmethod
    def from_str(cls, string: str) -> "Modifier": ...


@dataclass
class Modifier:
    """Base class for sound modifiers."""

    IDENTIFIER = ""

    @classmethod
    def from_str(cls, string: str) -> "Modifier":  # pylint: disable=unused-argument
        return Modifier()

    def as_str(self) -> str:
        return ""

    def modify_audio(self, audio: Audio) -> Audio:
        return audio

    def __eq__(self, other: object) -> bool:
        raise NotImplementedError


@dataclass
class SpeedChangeModifier(Modifier):
    """Modify speed of audio without changing pitch."""

    IDENTIFIER = "s"

    def __init__(self, speed_parameter: float) -> None:
        if speed_parameter <= 0.0:
            raise ModifierArgumentsInvalid
        self.speed_parameter = float(speed_parameter)

    @classmethod
    def from_str(cls, string: str) -> "SpeedChangeModifier":
        speed_parameter_str = string.strip(SpeedChangeModifier.IDENTIFIER)
        return SpeedChangeModifier(speed_parameter=float(speed_parameter_str))

    def as_str(self) -> str:
        return f"{self.IDENTIFIER}{self.speed_parameter}"

    def modify_audio(self, audio: Audio) -> Audio:
        stretched_data = librosa.effects.time_stretch(
            y=audio.data, rate=self.speed_parameter
        )
        return Audio(data=stretched_data, sample_rate=audio.sample_rate)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, SpeedChangeModifier):
            return False
        return self.speed_parameter == other.speed_parameter


@dataclass
class PitchChangeModifier(Modifier):
    """Modify pitch of audio without changing speed."""

    IDENTIFIER = "p"

    def __init__(self, pitch_parameter: float) -> None:
        self.pitch_parameter = float(pitch_parameter)

    @classmethod
    def from_str(cls, string: str) -> "PitchChangeModifier":
        pitch_parameter_str = string.strip(PitchChangeModifier.IDENTIFIER)
        return PitchChangeModifier(pitch_parameter=float(pitch_parameter_str))

    def as_str(self) -> str:
        return f"{self.IDENTIFIER}{self.pitch_parameter}"

    def modify_audio(self, audio: Audio) -> Audio:
        pitch_shifted_data = librosa.effects.pitch_shift(
            y=audio.data,
            sr=audio.sample_rate,
            n_steps=self.pitch_parameter,
        )
        return Audio(data=pitch_shifted_data, sample_rate=audio.sample_rate)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, PitchChangeModifier):
            return False
        return self.pitch_parameter == other.pitch_parameter


class AcceleratorDirection(enum.Enum):
    """Direction for accelerator modifier."""

    FORWARD = enum.auto()
    REVERSE = enum.auto()

    @classmethod
    def from_str(cls, value: str) -> "AcceleratorDirection":
        if value == "f":
            return AcceleratorDirection.FORWARD
        if value == "r":
            return AcceleratorDirection.REVERSE
        raise AcceleratorDirectionInvalid

    def to_str(self) -> str:
        if self == AcceleratorDirection.FORWARD:
            return "f"
        if self == AcceleratorDirection.REVERSE:
            return "r"
        raise AcceleratorDirectionInvalid


@dataclass
class AcceleratorModifier(Modifier):
    """Accelerate a single word with modifiable pitch and speed curves."""

    IDENTIFIER = "a"

    def __init__(
        self, speed_increment: float, count: int, direction: AcceleratorDirection
    ) -> None:
        if speed_increment <= 0.0:
            raise ModifierArgumentsInvalid
        if count <= 1.0:
            raise ModifierArgumentsInvalid

        self.speed_increment = speed_increment
        self.count = int(count)
        self.direction = direction

    @classmethod
    def from_str(cls, string: str) -> "AcceleratorModifier":
        parameter_str = string.strip(cls.IDENTIFIER)
        components = parameter_str.split("+")
        if len(components) < 2 or len(components) > 3:
            raise ModifierFormatNotCorrect

        speed_increment = float(components[0])
        count = int(components[1])
        direction = AcceleratorDirection.FORWARD
        if len(components) == 3:
            direction = AcceleratorDirection.from_str(components[2])

        return AcceleratorModifier(
            speed_increment=speed_increment, count=count, direction=direction
        )

    def as_str(self) -> str:
        return f"{self.IDENTIFIER}{self.speed_increment}+{self.count}+{self.direction.to_str()}"

    def modify_audio(self, audio: Audio) -> Audio:
        segment: Audio | None = None
        increments: Union[range, reversed]

        if self.direction == AcceleratorDirection.REVERSE:
            increments = reversed(range(0, self.count))
        else:
            increments = range(0, self.count)

        for increment in increments:
            if increment == 0.0:
                speed_increment = 1.0
            else:
                speed_increment = increment * self.speed_increment

            speed_modifier = SpeedChangeModifier(speed_parameter=speed_increment)
            if speed_modifier != 1.0:
                modified_audio = speed_modifier.modify_audio(audio=audio)
            else:
                modified_audio = copy.deepcopy(audio)

            if segment is None:
                segment = modified_audio
            else:
                segment += modified_audio

        if segment is None:
            raise ModifierArgumentsInvalid("No segments were created")

        return Audio(data=segment.data, sample_rate=segment.sample_rate)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, AcceleratorModifier):
            return False
        return (
            self.speed_increment == other.speed_increment
            and self.count == other.count
            and self.direction == other.direction
        )


@dataclass
class VolumeChangeModifier(Modifier):
    """Modify volume of audio."""

    IDENTIFIER = "v"

    def __init__(self, volume_parameter: float) -> None:
        self.volume_parameter = float(volume_parameter)

    @classmethod
    def from_str(cls, string: str) -> "VolumeChangeModifier":
        volume_parameter_str = string.strip(VolumeChangeModifier.IDENTIFIER)
        return VolumeChangeModifier(volume_parameter=float(volume_parameter_str))

    def as_str(self) -> str:
        return f"{self.IDENTIFIER}{self.volume_parameter}"

    def modify_audio(self, audio: Audio) -> Audio:
        return Audio(
            data=audio.data * self.volume_parameter, sample_rate=audio.sample_rate
        )

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, VolumeChangeModifier):
            return False
        return self.volume_parameter == other.volume_parameter


class CutModifier(Modifier):
    """Cut start and end of audio."""

    IDENTIFIER = "c"

    def __init__(self, start_time_s: float | None, end_time_s: float | None) -> None:
        if start_time_s is None and end_time_s is None:
            raise ModifierArgumentsInvalid
        if start_time_s is not None and start_time_s < 0:
            raise ModifierArgumentsInvalid
        if end_time_s is not None and end_time_s < 0:
            raise ModifierArgumentsInvalid
        if (
            start_time_s is not None
            and end_time_s is not None
            and start_time_s >= end_time_s
        ):
            raise ModifierArgumentsInvalid

        self.start_time_s = start_time_s
        self.end_time_s = end_time_s

    @classmethod
    def from_str(cls, string: str) -> "CutModifier":
        start_time_s_str, end_time_s_str = string.strip(CutModifier.IDENTIFIER).split(
            "-"
        )
        start_time_s = None if start_time_s_str == "" else float(start_time_s_str)
        end_time_s = None if end_time_s_str == "" else float(end_time_s_str)
        return CutModifier(start_time_s=start_time_s, end_time_s=end_time_s)

    def as_str(self) -> str:
        start_time_s_str = "" if self.start_time_s is None else f"{self.start_time_s}"
        end_time_s_str = "" if self.end_time_s is None else f"{self.end_time_s}"
        return f"{self.IDENTIFIER}{start_time_s_str}-{end_time_s_str}"

    def modify_audio(self, audio: Audio) -> Audio:
        start_index = (
            0
            if self.start_time_s is None
            else int(self.start_time_s * audio.sample_rate)
        )
        end_index = (
            len(audio.data)
            if self.end_time_s is None
            else int(self.end_time_s * audio.sample_rate)
        )
        return Audio(
            data=audio.data[start_index:end_index], sample_rate=audio.sample_rate
        )

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, CutModifier):
            return False
        return (
            self.start_time_s == other.start_time_s
            and self.end_time_s == other.end_time_s
        )


@dataclass
class NoOpModifier1(Modifier):
    """No-op modifier for testing."""

    IDENTIFIER = "z"

    def __init__(self, parameter: float) -> None:
        self.parameter = float(parameter)

    @classmethod
    def from_str(cls, string: str) -> "NoOpModifier1":
        parameter = string.strip(NoOpModifier1.IDENTIFIER)
        return NoOpModifier1(parameter=float(parameter))

    def as_str(self) -> str:
        return f"{self.IDENTIFIER}{self.parameter}"

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, NoOpModifier1):
            return False
        return self.parameter == other.parameter


@dataclass
class NoOpModifier2(Modifier):
    """No-op modifier for testing."""

    IDENTIFIER = "x"

    def __init__(self, parameter: float) -> None:
        self.parameter = float(parameter)

    @classmethod
    def from_str(cls, string: str) -> "NoOpModifier2":
        parameter = string.strip(NoOpModifier2.IDENTIFIER)
        return NoOpModifier2(parameter=float(parameter))

    def as_str(self) -> str:
        return f"{self.IDENTIFIER}{self.parameter}"

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, NoOpModifier2):
            return False
        return self.parameter == other.parameter


MODIFIERS: Dict[str, Type[ModifierClass]] = {
    CutModifier.IDENTIFIER: CutModifier,
    VolumeChangeModifier.IDENTIFIER: VolumeChangeModifier,
    SpeedChangeModifier.IDENTIFIER: SpeedChangeModifier,
    PitchChangeModifier.IDENTIFIER: PitchChangeModifier,
    AcceleratorModifier.IDENTIFIER: AcceleratorModifier,
    NoOpModifier1.IDENTIFIER: NoOpModifier1,
    NoOpModifier2.IDENTIFIER: NoOpModifier2,
}

Modifiers: TypeAlias = List[Modifier]
