import torch.multiprocessing as mp

mp.set_start_method("spawn", force=True)

import torch
from dataclasses import dataclass
import numpy as np
import os
from suno_utils.audio import Audio
from suno_utils.tasks.lyrics_alignment.shortest_path_aligner import (
    ShortestPathAligner,
    ShortestPathAlignerConfig,
)
from suno_utils.tasks.hoot import clean_text, get_cer, SAMPLE_RATE, Tokenizer

_process_local_aligner: ShortestPathAligner | None = None
_process_local_tokenizer: Tokenizer = None


@dataclass
class LoadAudioResult:
    # when audio_tensor is in shared memory, objects of this class are very cheap to transfer between processes
    audio_tensor: torch.Tensor
    waveform: list[float]


def _align_in_worker(prior_text: str, logits: torch.Tensor, **kwargs):
    assert _process_local_aligner is not None
    assert _process_local_tokenizer is not None
    logits_numpy = logits.detach().cpu().numpy()
    del logits

    alignment = _process_local_aligner.align(prior_text, logits_numpy, **kwargs)

    # calculate the cer
    basic_cleaned_lyrics = clean_text("".join([w["word"] for w in alignment]))
    hoot_lyrics = _process_local_tokenizer.decode_logits(
        logits_numpy, basic_cleaned_lyrics if basic_cleaned_lyrics.strip() else None
    )
    hoot_cer = get_cer(prior_text, hoot_lyrics)
    return alignment, hoot_lyrics, hoot_cer


def _load_audio_in_worker(path: str) -> LoadAudioResult:
    audio = Audio.from_file(path)

    # waveform chunks are .2 seconds by default
    CHUNK_SIZE_S = 0.2

    # calculate waveform RMS data
    audio_floats = audio.mono().array_float
    audio_floats_split = np.array_split(audio_floats, int(audio.duration_s / CHUNK_SIZE_S))
    audio_rms = [round(float(np.sqrt(np.mean(np.square(x)))), 5) for x in audio_floats_split]

    audio_tensor = torch.from_numpy(
        audio.convert(sample_rate=SAMPLE_RATE, byte_width=2, n_channels=1).array_float[None]
    )
    audio_tensor.share_memory_()

    return LoadAudioResult(
        audio_tensor=audio_tensor,
        waveform=audio_rms,
    )


def _init_worker(tokenizer_filepath: str, aligner_config: ShortestPathAlignerConfig):
    global _process_local_tokenizer
    global _process_local_aligner

    if tokenizer_filepath.startswith("s3://"):
        from suno_utils.utils.s3 import read_from_s3

        _process_local_tokenizer = read_from_s3(tokenizer_filepath, read_f=Tokenizer)
    else:
        _process_local_tokenizer = Tokenizer(tokenizer_filepath)

    _process_local_aligner = ShortestPathAligner.from_sentencepiece(
        _process_local_tokenizer._tokenizer,
        config=aligner_config,
    )

    # stick to our core
    torch.set_num_threads(1)
    os.environ["OMP_NUM_THREADS"] = "1"


class HootBackgroundWorker:
    def __init__(
        self, tokenizer_filepath: str, aligner_config: ShortestPathAlignerConfig, num_workers: int
    ):
        self.pool = mp.Pool(
            processes=num_workers,
            initializer=_init_worker,
            initargs=(tokenizer_filepath, aligner_config),
        )

    def cleanup(self):
        self.pool.close()
        self.pool.join()

    def align(self, prior_text: str, logits: torch.Tensor, **kwargs):
        # run shortest path aligner in a background process
        return self.pool.apply(_align_in_worker, (prior_text, logits), kwargs)

    def load_audio(self, path: str) -> LoadAudioResult:
        # load audio, resample for hoot, and calculate RMS in a background process
        return self.pool.apply(_load_audio_in_worker, (path,))