import re
import random
from torch.utils import data
from suno_utils.utils.text import read_jsonl
from suno_utils.utils.text import normalize_whitespace
from suno_utils.utils.lyrics import remove_speakers


def clean_text(text: str) -> str:
    """General text cleaning. A bit tight but makes the content very clean.

    Returns a cleaned text string that is expected to be recongizable by hoot.
    """
    text = "\n" + text
    text = text.replace("’", "'").lower()
    text = text.replace('"', "").lower()
    text = re.sub(r"\[.+?\]", " ", text)  # tags
    text = re.sub(r"\n.+?\:", " ", text)  # new line ends with :
    text = re.sub(r"\n.+?\：", " ", text)  # new line ends with :
    text = re.sub(r"\n\(.+?\)", " ", text)  # new line with ()
    text = re.sub(r"[\d]", " ", text)  # digits
    text = re.sub(r"▁", "", text)  # special stuff
    text = remove_speakers(text)
    text = re.sub(r"[^\w\'\s]", " ", text)  # keep only the words
    text = normalize_whitespace(text)
    return text


class SelfLyricSimDataset(data.Dataset):
    def __init__(self, split="train", input_length=300, num_samples=-1):
        assert split in ["train", "valid"]
        self.input_length = input_length
        self.num_samples = num_samples
        self.split = split

        # Load files
        self.filelist = self.get_split(split)
        print(
            f"{len(self.filelist)} files are available for self_lyric_sim {split} set"
        )

    def get_split(self, split):
        fl = read_jsonl(
            "/app/suno/minz/metadata/pairs_metadata.jsonl", max_lines=100000
        )
        random.seed(134)
        random.shuffle(fl)
        if split == "train":
            train_fl = fl[:-2000]
            if self.num_samples > 0:
                train_fl = train_fl[: self.num_samples]
            return train_fl
        elif split == "valid":
            valid_fl = fl[-2000:]
            if self.num_samples > 0:
                valid_fl = valid_fl[: self.num_samples]
            return valid_fl

    def __getitem__(self, index):
        # read data
        if self.split == "train":
            good_length = False
            while not good_length:
                lyrics = random.choice(self.filelist)["lyrics"]
                lyrics = clean_text(lyrics)
                if len(lyrics) > self.input_length:
                    good_length = True
        elif self.split == "valid":
            lyrics = self.filelist[index]["lyrics"]
            lyrics = clean_text(lyrics)
        if self.split == "train":
            start_ix_1 = random.randint(0, len(lyrics) - self.input_length - 1)
            lyrics_1 = lyrics[start_ix_1 : start_ix_1 + self.input_length]
            lyrics_1 = "[CLS]" + "[Lyrics]" + lyrics_1

            start_ix_2 = random.randint(0, len(lyrics) - self.input_length - 1)
            lyrics_2 = lyrics[start_ix_2 : start_ix_2 + self.input_length]
            lyrics_2 = "[CLS]" + "[Lyrics]" + lyrics_2

        elif self.split == "valid":
            # crop first and last tokens
            start_ix_1 = 0
            lyrics_1 = lyrics[start_ix_1 : start_ix_1 + self.input_length]
            lyrics_1 = "[CLS]" + "[Lyrics]" + lyrics_1

            start_ix_2 = max(0, len(lyrics) - self.input_length)
            lyrics_2 = lyrics[start_ix_2 : start_ix_2 + self.input_length]
            lyrics_2 = "[CLS]" + "[Lyrics]" + lyrics_2

        return lyrics_1, lyrics_2

    def __len__(self):
        if self.num_samples > 0:
            return self.num_samples
        else:
            return len(self.filelist)