import random
import string


class Perturbation:
    """Base class for SmoothLLM random perturbations."""

    def __init__(self, q):
        self.q = q
        self.alphabet = string.printable


class RandomSwapPerturbation(Perturbation):
    """Implementation of SmoothLLM random swap perturbations.
    See `RandomSwapPerturbation` in lines 1-5 of Algorithm 2."""

    def __init__(self, q):
        super().__init__(q)

    def __call__(self, s):
        list_s = list(s)
        sampled_indices = random.sample(range(len(s)), int(len(s) * self.q / 100))
        for i in sampled_indices:
            list_s[i] = random.choice(self.alphabet)
        return "".join(list_s)


class RandomPatchPerturbation(Perturbation):
    """Implementation of SmoothLLM random patch perturbations.
    See `RandomPatchPerturbation` in lines 6-10 of Algorithm 2."""

    def __init__(self, q):
        super().__init__(q)

    def __call__(self, s):
        list_s = list(s)
        substring_width = int(len(s) * self.q / 100)
        max_start = len(s) - substring_width
        start_index = random.randint(0, max_start)
        sampled_chars = "".join([random.choice(self.alphabet) for _ in range(substring_width)])
        list_s[start_index : start_index + substring_width] = sampled_chars
        return "".join(list_s)


class RandomInsertPerturbation(Perturbation):
    """Implementation of SmoothLLM random insert perturbations.
    See `RandomPatchPerturbation` in lines 11-17 of Algorithm 2."""

    def __init__(self, q):
        super().__init__(q)

    def __call__(self, s):
        list_s = list(s)
        sampled_indices = random.sample(range(len(s)), int(len(s) * self.q / 100))
        for i in sampled_indices:
            list_s.insert(i, random.choice(self.alphabet))
        return "".join(list_s)
