# THIS FILE IS DEPRECIATED, PLEASE USE modal_runner_chirp_v2_clips.py INSTEAD
import json
import time
from typing import Any
import pathlib
from uuid import uuid4

import modal
from torchaudio.pipelines import HDEMUCS_HIGH_MUSDB_PLUS
from openai import OpenAI


from suno_utils.audio import Audio
from suno_utils.gpt import chirp_v2
from suno_utils.worker.generate_song_lyrics import (
    generate_song_lyrics_with_genre_tags,
    moderate_user_inputs,
    ModerationSuccess,
    ModerationError,
)

from suno_utils.worker.loader import S3Loader
from suno_utils.worker.modal_base import MODAL_MOUNTS
from suno_utils.worker.schema import QueueItem
from suno_utils.worker.utils import recursive_ls_dir


base_image = (
    modal.Image.debian_slim()
    .apt_install("curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3")
    .run_commands(
        [
            'curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"',
            "unzip -q awscliv2.zip",
            "./aws/install",
        ]
    )
    .pip_install(
        "torch==2.1.0.dev20230831+cu118",
        "torchaudio==2.2.0.dev20230831+cu118",
        index_url="https://download.pytorch.org/whl/nightly/cu118",
    )
    .pip_install_private_repos(
        "github.com/suno-ai/glockenspiel.git@f05e2f251#subdirectory=descript-audio-codec&egg=descript-audio-codec",
        git_user="mcamac",
        secrets=[modal.Secret.from_name("victor-modal-github-token")],
    )
    .pip_install(
        "boto3",
        "transformers",
        "tokenizers",
        "encodec",
        "ctc_segmentation",
        "psutil",
        "redis",
        "gradio",
        "pydantic",
        "nnAudio",
    )
    .pip_install_from_pyproject(
        str(pathlib.Path(__file__).parent.parent.parent / "pyproject.toml"),
    )
    .run_commands(
        "FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation",
    )
)


MOUNT_PATH = "/suno/models"


class ChirpV2Worker(S3Loader):
    gpu_id: int

    def __init__(self, gpu_id, bg_image=None):
        super().__init__(bg_image=bg_image)

        self.gpu_id = gpu_id
        self.bg_image = bg_image

    def preload(self):
        start_time = time.time()
        chirp_v2.preload_models(cache_dir=MOUNT_PATH, local_whisper="/suno/models/whisper")

        finish_time = time.time()
        print(f"Preloading took {finish_time - start_time}s")

    @staticmethod
    def download_models(dir_path=MOUNT_PATH):
        """Use AWS CLI to download models if they don't exist."""
        chirp_v2.preload_models(fetch_only=True, cache_dir=dir_path)
        recursive_ls_dir(dir_path)

        HDEMUCS_HIGH_MUSDB_PLUS.get_model()

        import whisper

        whisper.load_model("small")
        whisper.load_model("small.en")

    def process_item(self, item: QueueItem) -> tuple[list[Audio], list[Any]]:
        history_audio = self._load_audio_prompt(item) if item.prompt_audio else None

        options = item.metadata.get("options", {}) or {}

        n_batch = 2
        chaos = options.get("chaos", 1)
        chaos = float(chaos)

        text = item.prompt_text or ""
        text_tags = item.metadata.get("tags", None)
        if text_tags == "random":
            text_tags = None
        detected_lang = chirp_v2._get_text_lang(text)
        cfg_coef = 1.25
        # cfg_coef = float(options.get("lyrics_strength", 1.2))
        # condition on non-english, higher cfg text
        if (detected_lang is not None and not detected_lang == "en") and text:
            cfg_coef = 1.7
            print(f"detected non-english lyrics, increase cfg, lang is {detected_lang}")
        # condition on genre, rap
        elif text_tags and ("rap" in text_tags or "hip-hop" in text_tags):
            cfg_coef = 1.4
        text_neg_tags = None
        if "instrumental" not in (text_tags or "") and len(text.split()) > 5:
            text_neg_tags = "instrumental noise"
        # cfg_coef_tags = float(options.get("tags_strength", 1.75))
        cfg = chirp_v2.GenerationConfig(
            text=text,
            text_tags=text_tags or None,
            text_neg_tags=text_neg_tags,
            n_batch=n_batch,
            max_gen_duration_s=80,
            cfg_coef=cfg_coef,
            cfg_coef_tags=1.9,
            temp_semantic=0.9,
            temp_coarse=0.85,
            min_text_offset=128,
            top_k_semantic=1000,
            top_k_coarse=1000,
            top_p_semantic=None,
            top_p_coarse=None,
            rep_penality=0,
        )

        MAX_HISTORY_DURATION_S = 50
        MAX_CONTINUATION_DURATION_S = 60
        if isinstance(history_audio, Audio):
            history_audio = history_audio.get_segment(
                from_s=max(0, history_audio.duration_s - MAX_HISTORY_DURATION_S)
            )
            cfg = cfg.modify(max_gen_duration_s=MAX_CONTINUATION_DURATION_S)
        elif history_audio is not None:  # NPZ
            # history_arrs are passed through GenerationConfig
            cfg = cfg.modify(
                max_gen_duration_s=MAX_CONTINUATION_DURATION_S,
                history_arr=history_audio,
                cfg_coef_tags=0,
            )
            history_audio = None

        _, raw_arrays, audios = chirp_v2.generate_audio(
            cfg,
            chaos_level=chaos,
            history_audio=history_audio,
            return_raw_arrays=True,
            max_history_duration_s=MAX_HISTORY_DURATION_S,
        )

        return audios, raw_arrays


aws_secret = modal.Secret.from_name("studio-aws")
SECRETS = [
    aws_secret,
    modal.Secret.from_dict(
        {
            "SUNO_ASSETS_PATH": "/suno/models/assets",
            "XDG_CACHE_HOME": "/suno/models/",
        }
    ),
    modal.Secret.from_name("openai-secret"),
]


def download_model_wrapper_a4():
    ChirpV2Worker.download_models()


image = base_image.run_function(download_model_wrapper_a4, secrets=SECRETS)

ENV_NAME = "staging"
STUB_NAME = f"chirp-v2-{ENV_NAME}"
stub = modal.Stub(STUB_NAME, image=image)


CONCURRENCY_LIMITS = {
    "priority": 80,
    "staging": 600,
}


@stub.cls(
    cpu=2.0,
    # memory=16384,
    gpu=modal.gpu.A10G(count=1),
    secrets=SECRETS,
    timeout=300,
    container_idle_timeout=400,
    mounts=MODAL_MOUNTS,
    memory=60000,
    retries=modal.Retries(
        max_retries=3,
        backoff_coefficient=2.0,
        initial_delay=10.0,
    ),
    keep_warm=1,
    concurrency_limit=CONCURRENCY_LIMITS.get(ENV_NAME, 80),
)
class ChirpV2Stub:
    def __enter__(self):
        import torch

        num_gpus = torch.cuda.device_count()
        print(f"Found {num_gpus} GPUs.")
        recursive_ls_dir("/suno/models")

        self.worker = ChirpV2Worker(0, bg_image="/suno/models/assets/wave-bg-2.png")
        self.worker.preload()

        self.openai_client = OpenAI()

    @modal.method()
    def generate(self, queue_item: str):
        import json

        from suno_utils.worker.schema import QueueItem

        print(queue_item)
        start_time = time.time()
        item = QueueItem(**json.loads(queue_item))
        item_id = item.id
        print(item.metadata)
        is_square = True
        ids = []

        if user_prompt := item.metadata.get("gpt_description_prompt"):
            try:
                title, lyrics, genre_tags = generate_song_lyrics_with_genre_tags(
                    self.openai_client, chirp_v2.text_lang_model, user_prompt
                )
            except ModerationError as e:
                import traceback

                self.worker.notify_finish(
                    item,
                    {
                        "id": item.id,
                        "type": "error",
                        "error_message": str(e),
                    },
                )
                traceback.print_exc()
                return

            item.prompt_text = lyrics
            item.metadata["tags"] = " ".join(genre_tags)
            print("generated", item.metadata["tags"], lyrics)

            self.worker.notify_progress(
                item,
                {
                    "id": item.id,
                    "type": "lyrics",
                    "title": title,
                    "text": item.prompt_text,
                    "tags": item.metadata["tags"],
                },
            )
        else:
            tags = item.metadata.get("tags") or ""
            moderation_result = moderate_user_inputs(self.openai_client, item.prompt_text or "", tags)
            print(moderation_result)
            if not isinstance(moderation_result, ModerationSuccess):
                self.worker.notify_finish(
                    item,
                    {
                        "id": item.id,
                        "type": "error",
                        "error_type": "moderation_failure",
                        "error_message": moderation_result.err_msg,
                    },
                )
                return

        f = modal.Function.lookup("stable-diffusion-xl-beta-3", "StableDiffusion.generate_image")
        fn_call = f.spawn(
            item.prompt_text or "",
            item_id,
            num_images=1,
            tags=item.metadata.get("tags", None),
        )

        try:
            audios, raw_arrays = self.worker.process_item(item)
            items = []
            for i, audio in enumerate(audios):
                new_id = f"{item_id}_{i}"
                ids.append(new_id)
                item.id = new_id
                # item.prompt_text = trimmed[i]
                items.append(item.json())
                self.worker._write_audio_only(
                    item,
                    audio,
                )

                self.worker._write_npz(
                    item,
                    raw_arrays[i],
                )

            image_url = fn_call.get(timeout=None)
            print("image located", image_url)

            self.worker.notify_finish(
                item,
                {
                    "id": item_id,
                    "model": "chirp-v2-test-xxl",
                    "n_audios": len(ids),
                    "ok": 1,
                    "gen_duration": time.time() - start_time,
                    "ids": ids,
                    "durations": [audio.duration_s for audio in audios],
                },
                queue_name="results:q",
            )

            video_time = time.time()
            f = modal.Function.lookup("dummy-v2-staging", "DummyV0Stub.write_video")
            # Start the video jobs.
            for item in items:
                f.spawn(item, image_url.split("/")[-1], is_square)

            print("Spawning videos", time.time() - video_time)
            return item_id

        except Exception as e:
            import traceback

            finish_time = time.time()
            self.worker.notify_finish(
                item,
                {
                    "id": item_id,
                    "model": "chirp-v2-test-xxl",
                    "n_audios": len(ids),
                    "ok": 0,
                    "gen_duration": finish_time - start_time,
                },
                queue_name="results:q",
            )

            traceback.print_exc()

            raise e

        return item.id


@stub.local_entrypoint()
def main():
    uid = str(uuid4())
    inputs = [
        json.dumps(
            dict(
                id=uid,
                prompt_text="""Hello world""",
                metadata={"tags": "pop"},
            ),
        ),
        json.dumps(
            dict(
                id=uid,
                prompt_text="""你好""",
                metadata={"tags": "pop"},
            ),
        ),
    ]

    model = ChirpV2Stub()
    print(uid)
    for input in inputs:
        model.generate.remote(
            input,
        )