"""Audio stem (diffusion) application on modal.

Code --> Diffusion --> Audio
"""

import json
import modal
import time
import os
import asyncio
import torch
from uuid import uuid4
from typing import Optional, Dict, Tuple
from datadog import initialize, statsd
import numpy as np
from concurrent.futures import ThreadPoolExecutor, Future
import threading
import queue

from suno_utils.worker.loader import S3Loader
from suno_utils.worker.schema import QueueItem, HistoryPrompt
from suno_utils.audio import Audio
from suno_utils.worker.settings import s3_client
from suno_utils.worker.utils import retry_decorator
from suno_utils.worker.modal_model_configs import MODEL_CONFIG_DICT
from suno_utils.diffusion_stems.generation import INSTRUMENT_CATEGORIES, STEM_TYPE_ID_TO_NAME
import suno_utils.diffusion.generation as diffusion_gen
from suno_utils.worker.tracing import distributed_trace, tracer
from suno_utils.worker.modal_base import get_modal_base_image_with_flash_attention
from suno_utils.worker.modal_model_volume import MODEL_STORE_VOLUME_DIR, model_store_volume
from suno_utils.gpt.rpc_zmq_upsample import start_service_processes as start_upsample_service_processes
from suno_utils.tasks.upsample_engine import DiffusionGenerationConfig, Request
from suno_utils.tasks.dac_vae_fixed_25hz import (
    decode_stream_to_full_audio,
    encode_overlap,
    preload_models,
)
############## CHANGE THESE ##############

DEPLOYMENT_TYPE = "dev"  # dev, prod
MODEL_CONFIG = MODEL_CONFIG_DICT[
    "stems_v1_12_output"
]  # stems_v0, stems_v1, stems_v1_8_output, stems_v1_12_output

##########################################

fast_retry_s3_download = retry_decorator(3, wait_seconds=1)(s3_client.download_fileobj)
MOUNT_PATH = "/tmp/suno/models"
MODEL = MODEL_CONFIG.model
MODEL_GAIN_ADJUST = MODEL_CONFIG.gain_adjust
DIFFUSION_CKPT_PATH = MODEL_CONFIG.diffusion_ckpt_path
# print(f"Stem model: {DIFFUSION_CKPT_PATH}")
aws_secret = modal.Secret.from_name("studio-aws")
SECRETS = [
    aws_secret,
    modal.Secret.from_name("api-callback-token"),
    modal.Secret.from_dict(
        {
            "DD_SITE": "datadoghq.com",
            "DD_ENV": DEPLOYMENT_TYPE,
            "DD_SERVICE": "stem-worker",
            "DD_LOGS_ENABLED": "false",
            "DD_TRACE_ENABLED": "true",
        },
    ),
    modal.Secret.from_name("datadog-metrics"),
]

# These queue names must correspond to the ones in streaming_api ENV_CONFIGS
CHUNK_QUEUE_NAME = f"chunk-queue-{DEPLOYMENT_TYPE}"
STREAM_KEY_QUEUE_NAME = f"stream-key-queue-{DEPLOYMENT_TYPE}"
UPSAMPLE_DDOG_SERVICE = "stem-worker"

audio_chunk_queue = modal.Queue.from_name(CHUNK_QUEUE_NAME, create_if_missing=True)
# Queue used by streaming_api to determine which ephemeral stream_key to stream from.
# This corresponds to stream_key_queue_name in streaming_api.
appstream_key_queue = modal.Queue.from_name(STREAM_KEY_QUEUE_NAME, create_if_missing=True)


base_image = get_modal_base_image_with_flash_attention()


class StemWorker(S3Loader):
    def __init__(self):
        S3Loader.__init__(self)
        # configure the stem's min chunk size here
        start_time = time.time()
        print("Start loading models")

        tokenizer_path = diffusion_gen.get_model_if_needed(
            diffusion_gen.TOKENIZER_FILEPATH,
            cache_dir=MOUNT_PATH,
        )

        codec_path = diffusion_gen.get_model_if_needed(
            MODEL_CONFIG.codec_ckpt_path, cache_dir=MOUNT_PATH
        )
        preload_models(codec_path)

        print(f"Dit model path: {DIFFUSION_CKPT_PATH}")
        dit_model_path = diffusion_gen.get_model_if_needed(DIFFUSION_CKPT_PATH, cache_dir=MOUNT_PATH)

        self.engine_rpc_client = start_upsample_service_processes(
            min_chunk_size=25 * (60 if MODEL == "stems_v1" else 30),
            compile=True,
            dit_model_filepath=dit_model_path,
            tokenizer_filepath=tokenizer_path,
        )
        self.model_config = self.engine_rpc_client.get_model_config()


def download_models_to_modal_disk6(dir_path=MOUNT_PATH):
    """Download diffusion models."""
    print("Start downloading models")
    _ = diffusion_gen.get_model_if_needed(diffusion_gen.CODEC_FILEPATH, cache_dir=dir_path)
    _ = diffusion_gen.get_model_if_needed(DIFFUSION_CKPT_PATH, cache_dir=dir_path)
    print("Finish downloading models")


image = base_image.pip_install("transformers==4.44.0").add_local_python_source("suno_utils", copy=False)
APP_NAME = f"stems-{MODEL_CONFIG.model}-{DEPLOYMENT_TYPE}"
app = modal.App(APP_NAME, image=image)


@app.cls(
    gpu=["H100"],  # H100 is faster, but if limited by resources, burst into A100
    cpu=4,
    secrets=SECRETS,
    timeout=240,  # the wav encoding can take a while
    scaledown_window=240,
    memory=15000,
    retries=modal.Retries(
        max_retries=1,
        backoff_coefficient=2.0,
        initial_delay=5.0,
    ),
    min_containers=2 if DEPLOYMENT_TYPE == "dev" else 2,
    max_containers=4 if DEPLOYMENT_TYPE == "dev" else 500,
    buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 2,
    volumes={MODEL_STORE_VOLUME_DIR: model_store_volume},
)
@modal.concurrent(max_inputs=2)
class StemStub:
    def __init__(self):
        import torch

        num_gpus = torch.cuda.device_count()
        device_name = torch.cuda.get_device_name()
        print(f"Found {num_gpus} GPUs. Using CUDA device: {device_name}")
        self.worker = StemWorker()
        self.modal_f_cycle = modal.Cls.lookup(
            f"cycle-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}",
            "CycleStub",
        )().encode_audio

        options = {"statsd_host": "127.0.0.1", "statsd_port": 8125}

        initialize(**options)

        # Separate thread pools for different workloads
        self.download_executor = ThreadPoolExecutor(
            max_workers=8,  # Scale download workers independently
            thread_name_prefix="audio_download",
        )
        self.encode_executor = ThreadPoolExecutor(
            max_workers=4,  # CPU encoding workers
            thread_name_prefix="audio_encode",
        )

        # Pipeline queues
        self.download_queue = queue.Queue(maxsize=20)  # Raw audio queue
        self.ready_queue = queue.Queue(maxsize=10)  # Processed audio ready for GPU

        # Request tracking
        self.pending_requests: Dict[str, Future] = {}
        self.request_lock = threading.Lock()

        # Start background workers
        self._start_pipeline_workers()

    def _start_pipeline_workers(self):
        """Start background workers for the audio processing pipeline."""

        def download_worker():
            """Downloads audio files from S3."""
            while True:
                try:
                    item_id, audio_s3_path, history = self.download_queue.get(timeout=30)

                    try:
                        if history is not None and history.history_latents is not None:
                            raw_audio = history.history_latents
                            duration = round(raw_audio.shape[0] / 25)
                            # Skip encoding for latents
                            self.ready_queue.put((item_id, raw_audio, duration, True))  # is_latent=True
                        else:
                            print(f"Downloading audio for {item_id}: {audio_s3_path}")
                            raw_audio = Audio.from_s3(audio_s3_path, n_channels=2)
                            duration = round(raw_audio.duration_s, 2)
                            # Submit for encoding
                            self.encode_executor.submit(
                                self._encode_worker, item_id, raw_audio, duration
                            )

                    except Exception as e:
                        print(f"Download failed for {item_id}: {e}")
                        with self.request_lock:
                            if item_id in self.pending_requests:
                                self.pending_requests[item_id].set_exception(e)

                except queue.Empty:
                    continue

        def encode_worker_wrapper():
            """Manages encoding workers."""
            # The actual encoding happens in _encode_worker submitted to encode_executor
            pass

        # Start download workers (multiple threads)
        for i in range(4):  # Multiple download threads
            thread = threading.Thread(target=download_worker, daemon=True, name=f"download_worker_{i}")
            thread.start()

    def _encode_worker(self, item_id: str, raw_audio: Audio, duration: float):
        """Encodes raw audio (runs in encode_executor)."""
        try:
            print(f"Encoding audio for {item_id}")
            encoded_audio = encode_overlap(raw_audio, normalize_volume=False)
            self.ready_queue.put((item_id, encoded_audio, duration, False))  # is_latent=False
        except Exception as e:
            print(f"Encoding failed for {item_id}: {e}")
            with self.request_lock:
                if item_id in self.pending_requests:
                    self.pending_requests[item_id].set_exception(e)

    def _submit_audio_request(self, item: QueueItem, history: Optional[HistoryPrompt]) -> Future:
        """Submit audio loading request and return future."""
        future = Future()

        with self.request_lock:
            self.pending_requests[item.id] = future

        # Determine audio path
        audio_s3_path = item.prompt_audio
        if not audio_s3_path.startswith("s3://"):
            audio_s3_path = f"s3://suno-data-uploads/studio/uploads/{audio_s3_path}.mp3"

        # Submit to download queue
        try:
            self.download_queue.put((item.id, audio_s3_path, history), timeout=1)
        except queue.Full:
            future.set_exception(Exception(f"Download queue full for {item.id}"))

        return future

    def _get_ready_audio(self, item_id: str, timeout: float = 30) -> Tuple[np.ndarray, float]:
        """Get processed audio that's ready for GPU."""
        deadline = time.time() + timeout

        while time.time() < deadline:
            try:
                ready_item_id, audio, duration, is_latent = self.ready_queue.get(timeout=1)

                if ready_item_id == item_id:
                    # Clean up tracking
                    with self.request_lock:
                        self.pending_requests.pop(item_id, None)
                    return audio, duration
                else:
                    # Put back if it's for a different request
                    self.ready_queue.put((ready_item_id, audio, duration, is_latent))

            except queue.Empty:
                continue

        raise TimeoutError(f"Audio not ready for {item_id} within {timeout}s")

    @modal.method()
    @distributed_trace("stem", UPSAMPLE_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE)
    def stem(self, queue_item: str, history: Optional[HistoryPrompt] = None) -> None:
        """Main stem processing method with pipelined audio loading."""
        tracer.current_span().set_tag("deployment_type", DEPLOYMENT_TYPE)
        item = QueueItem(**json.loads(queue_item))
        ids = item.multi_ids
        assert ids is not None
        device_name = torch.cuda.get_device_name()
        dd_tags = [
            f"model:{MODEL}",
            f"env:{DEPLOYMENT_TYPE}",
            f"env_name:{APP_NAME}",
            f"modal_cloud_provider:{os.environ.get('MODAL_CLOUD_PROVIDER', 'unknown')}",
            f"modal_environment:{os.environ.get('MODAL_ENVIRONMENT', 'unknown')}",
            f"modal_image_id:{os.environ.get('MODAL_IMAGE_ID', 'unknown')}",
            f"modal_region:{os.environ.get('MODAL_REGION', 'unknown')}",
            f"gen_type:{item.metadata.get('type', 'unknown')}",
            f"gen_task:{item.metadata.get('task', 'unknown')}",
            f"cuda_device:{device_name}",
        ]

        start_time = time.time()

        # Submit audio loading immediately (non-blocking)
        audio_future = self._submit_audio_request(item, history)

        # Each unique Chirp request is assigned a unique ephemeral stream_key.
        # Replace item in the stream_key_queue.
        # streaming_api streams from the first item in the stream_key_queue.
        stream_key = "stream-key-" + str(uuid4())
        appstream_key_queue.get(block=False, partition=item.id)
        appstream_key_queue.put(stream_key, partition=item.id, partition_ttl=600)
        # If specified, use stream_key to read tokens from the token queue.
        # Otherwise, fall back to the item id.
        partition = stream_key or item.id
        print(f"Got {item}")

        if MODEL == "stems_v1":
            stem_type_ids = item.metadata.get("stem_type_id", 0)
            if not isinstance(stem_type_ids, list):
                stem_type_ids = [stem_type_ids]
            stem_type_names = [STEM_TYPE_ID_TO_NAME[str(stem_type_id)] for stem_type_id in stem_type_ids]

            # if stem_type_group_name is set, use it as the stem type name
            stem_type_group_name = item.metadata.get("stem_type_group_name")
            if stem_type_group_name:
                assert stem_type_group_name in INSTRUMENT_CATEGORIES, stem_type_group_name
                stem_type_names = [stem_type_group_name]

            stem_type_cfg_scale = item.metadata.get("stem_type_cfg_scale", 3)

            stem_task = item.metadata.get("stem_task", "extract")
            if stem_task is None:
                stem_task = "extract"
            print(
                f"{item.id} has stem type ids "
                f"{stem_type_ids}, stem type names "
                f"{stem_type_names}, stem type cfg scale "
                f"{stem_type_cfg_scale}."
            )
            tags = f"{stem_task}"
            if stem_task in ["add", "remove", "extract"]:
                tags += f" [{', '.join(stem_type_names)}]"
        elif MODEL == "stems_v1_8_output" or MODEL == "stems_v1_12_output":
            tags = "extract [split_karaoke]"
            stem_type_cfg_scale = item.metadata.get("stem_type_cfg_scale", 1)
            if stem_type_cfg_scale != 1:
                print(
                    f"WARNING: Stem {item.id}: stem_type_cfg_scale is not 1: {stem_type_cfg_scale}. This is much slower and probably doesnt improve quality"
                )
        else:
            raise ValueError(f"Unknown model: {MODEL}")
        print(f"Stem {item.id}: tags: {tags}")

        diffusion_inference_config = {
            "lyrics": tags,
            "text_cfg_coef": stem_type_cfg_scale,
            "steps": 4,
            "codec_scale_factor": 0.4,
            "scale_ctx_vector": 0.4,
            "noise_ctx_level": 0.05,
        }
        if item.is_diff_infill:
            assert history is not None
            history_latents = torch.from_numpy(history.history_latents)
            future_latents = torch.from_numpy(history.future_latents)
            diffusion_inference_config["infill_prefix_latents"] = history_latents
            diffusion_inference_config["infill_suffix_latents"] = future_latents

        # Note that experiment config will overwrite default config
        if special_config := item.metadata.get("model_config"):
            # has specific config instruction
            print(f"Stem {item.id}: load specific config: {special_config}.")
            # update the custom config
            # note this will overwrite existing parameters
            diffusion_inference_config.update(**special_config)
        # Note forced infer config will overwrite experiment config
        if special_config := item.metadata.get("forced_infer_config"):
            # has specific config instruction
            print(f"Stem {item.id}: load forced infer config: {special_config}")
            # update the custom config
            # note this will overwrite existing parameters
            diffusion_inference_config.update(
                **{k: v for k, v in special_config.items() if v is not None}
            )
        # Filter kwargs that are valid properties of GenerationConfig
        valid_gconf_props = set(vars(DiffusionGenerationConfig()).keys())
        invalid_keys = set(diffusion_inference_config.keys()) - valid_gconf_props
        if invalid_keys:
            print(f"WARNING: Stem {item.id}: Invalid config keys found: {invalid_keys}")
        diffusion_inference_config = {
            k: v for k, v in diffusion_inference_config.items() if k in valid_gconf_props
        }
        print(f"Stem {item.id}: diffusion_inference_config: {diffusion_inference_config}")

        setup_time = time.time() - start_time
        print(f"Setup for {item.id} completed in {setup_time:.2f}s, waiting for audio...")

        # Now get the processed audio (should be ready or nearly ready)
        try:
            audio, expected_duration = self._get_ready_audio(item.id, timeout=30)
        except Exception as e:
            print(f"Failed to get audio for {item.id}: {e}")
            raise

        audio_ready_time = time.time() - start_time
        print(f"Audio ready for {item.id} at {audio_ready_time:.2f}s (setup took {setup_time:.2f}s)")
        print(
            f"{item.id} got parent audio. Expected duration {expected_duration}s, {audio_ready_time:.2f}s elapsed."
        )

        if expected_duration > 16 * 60:  # 20 MB modal queue size
            raise ValueError(f"{item.id} Parent audio is too long: {expected_duration}s")

        if item.is_diff_infill:
            assert audio.shape[0] <= self.worker.model_config["stem_ctx_len"]

        diffusion_generation_config = DiffusionGenerationConfig(
            audio=audio, drop_semantic_tokens=True, **diffusion_inference_config
        )
        print(f"Stem {item.id} adding request")
        print(audio.shape)
        with tracer.trace("generate"):
            pad_token = self.worker.model_config["cond_semantic_n_vocab"] - 1
            job = self.worker.engine_rpc_client.add_request(
                Request(
                    item.id,
                    diffusion_generation_config,
                    input_tokens_finished=True,
                    tokens=np.full((audio.shape[0]), pad_token),
                    stem_ctx_latents=audio,
                )
            )
            while not self.worker.engine_rpc_client.get_job_state(job)["completed"]:
                time.sleep(0.05)
            latents = self.worker.engine_rpc_client.get_generated_codes(job)
        print(f"Stem {item.id} got latents: {latents.shape}")

        is_vocal_only = (
            MODEL == "stems_v1_12_output" and len(ids) == 2
        )  # if only 2 ids, then we are doing vocal+complement only
        if is_vocal_only:
            print(f"Stem {item.id} is vocal only")
        else:
            assert len(ids) == latents.shape[1], f"{len(ids)} != {latents.shape[1]}"

        # prepare outputs with (stem_queue_item, latent, out_audio)
        outputs = []
        outputs_to_write_audio = []
        for i in range(latents.shape[1]):
            if is_vocal_only:
                stem_id = f"{ids[0]}_stem_{i}"  # we use the first id for the stems
            else:
                stem_id = ids[i]
            stem_queue_item = item.model_copy(update={"id": stem_id})
            latent = latents[:, i]
            stride = 25 * 10
            if audio.shape[0] < stride:
                print(f"Stem {item.id} audio is too short, using stride {stride}")
                stride = 25
            out_audio = decode_stream_to_full_audio(
                latent,
                n_stride_tokens=stride,
            )  # TODO: this is pretty slow, should batch or use larger stride
            print(f"Stem {item.id} got audio: {out_audio.duration_s}")
            outputs.append((stem_queue_item, latent, out_audio))

        # prepare vocal and complement audio
        if is_vocal_only:
            print(f"Stem {item.id} is vocal only, preparing vocal and complement audio")
            # first 2 outputs are lead and backing vocals
            vocal_audio = Audio.sum([o[-1] for o in outputs[:2]])  # sum vocals and backing
            complement_audio = Audio.sum([o[-1] for o in outputs[2:]])  # sum the rest
            vocal_latent = encode_overlap(vocal_audio)
            complement_latent = encode_overlap(complement_audio)
            outputs.append((item.model_copy(update={"id": ids[0]}), vocal_latent, vocal_audio))
            outputs.append((item.model_copy(update={"id": ids[1]}), complement_latent, complement_audio))
            print(
                f"Stem {item.id} got vocal and complement audio: {vocal_audio.duration_s} and {complement_audio.duration_s}"
            )
            outputs_to_write_audio = outputs[-2:]  # only write vocal and complement
        else:
            outputs_to_write_audio = outputs

        # write empty bytes to modal to signal end of stream
        audio_chunk_queue.put(
            b"",
            partition=partition,
            partition_ttl=600,
            block=False,
        )
        if (gen_request_start_time := item.metadata.get("gen_request_start_time", None)) is not None:
            first_audio_write_time_millis = int(time.time() * 1000)
            elapsed_millis = first_audio_write_time_millis - gen_request_start_time
            if elapsed_millis > 0:
                print(
                    f"StemStub ({item.id}): Finish audio write time from gen request: {elapsed_millis // 1000}s"
                )
                if elapsed_millis > 600_000:
                    print(
                        f"StemStub ({item.id}): Finish audio write time from gen request took >10m: {elapsed_millis // 1000}s"
                    )
                statsd.distribution(
                    "decoder.audio_write_finish_millis.distribution",
                    elapsed_millis,
                    tags=dd_tags,
                )
        print(f"Stem {item.id}: {(time.time() - start_time):.2f}s elapsed.")

        with ThreadPoolExecutor(max_workers=2) as executor:
            # save npz for the vae latents
            print(f"Stem {item.id}: latents shape: {latents.shape}")
            npz_futures = [
                executor.submit(
                    self.worker._write_vae_latents_npz,
                    x_queue_item,
                    x,
                    MODEL_CONFIG.vae_version.value,
                    APP_NAME,
                )
                for x_queue_item, x, _ in outputs
            ]

            # save audio
            audio_t0 = time.time()

            async def _write_all_audio():
                await asyncio.gather(
                    *[
                        self.worker._write_audio_only_async(
                            stem_queue_item, stem_audio, gain_adjust=MODEL_GAIN_ADJUST, write_wav=True
                        )
                        for stem_queue_item, _, stem_audio in outputs_to_write_audio
                    ],
                )

            asyncio.run(_write_all_audio())
            print(f"Stem {item.id} audio encode and upload: {(time.time() - audio_t0):.2f}s")

            print(
                f"{item.id} uploaded stem and complement audio., Took {round(time.time() - start_time, 2)} seconds"
            )

            for future in npz_futures:
                future.result()

            print(f"Stem {item.id} vae latents npz: {(time.time() - start_time):.2f}s elapsed")

        # Clear multi_ids to prevent n^2 notifications
        # TODO clarify intent of multi_ids and fix this
        del item.metadata["multi_ids"]

        for index, id in enumerate(ids):
            stem_audio = outputs_to_write_audio[index][-1]
            stem_is_silent = bool(stem_audio.loudness < -45)
            item.notify_progress(
                {
                    "id": id,
                    "model": item.model_name,
                    "n_audios": 1,
                    "ok": True,
                    "gen_duration": time.time() - start_time,
                    "ids": [id],
                    "durations": [round(stem_audio.duration_s, 3)],
                    "is_silent": stem_is_silent,
                    **({"first_in_multi_bundle": True} if index == 0 else {}),
                },
            )

        self.worker.engine_rpc_client.remove_job(job)

        for id in ids:
            _ = self.modal_f_cycle.spawn(
                f"s3://suno-data-uploads/studio/uploads/{id}.mp3",
                s3_npz_id=id,
                encode_vae_version=None,
            )
        print(f"Done with stem {item.id}")

    @modal.exit()
    def cleanup_processes(self):
        print("StemStub: Cleaning up processes")
        self.download_executor.shutdown(wait=True)
        self.encode_executor.shutdown(wait=True)
        self.worker.engine_rpc_client.terminate_processes()


def test_stems_v1_8_output():
    model = StemStub()

    id = str(uuid4())
    output_ids = [f"{id}_{i}" for i in range(8)]
    input = {
        "id": id,
        "prompt_audio": "a5e2198a-f352-4abb-9a24-7f81b143ded3",  # stone
        "model_name": "stems",
        "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "multi_ids": output_ids,
        },
    }
    model.stem.remote(json.dumps(input))


def test_stems_v1_12_output():
    model = StemStub()

    id = str(uuid4())
    output_ids = [f"{id}_{i}" for i in range(12)]
    input = {
        "id": id,
        "prompt_audio": "a5e2198a-f352-4abb-9a24-7f81b143ded3",  # stone
        "model_name": "stems",
        "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "multi_ids": output_ids,
        },
    }
    model.stem.remote(json.dumps(input))

    print("Done with stems_v1_12_output")

    print("Testing vocal only")
    # test vocal only
    id = str(uuid4())
    input = {
        "id": id,
        "prompt_audio": "a5e2198a-f352-4abb-9a24-7f81b143ded3",  # stone
        "model_name": "stems",
        "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "multi_ids": [str(uuid4()) + "_vocal", str(uuid4()) + "_complement"],
        },
    }
    model.stem.remote(json.dumps(input))


def stress_test_stems_v1_12_output():
    model = StemStub()

    for i in range(10000):
        id = str(uuid4())
        input = {
            "id": id,
            "prompt_audio": "a5e2198a-f352-4abb-9a24-7f81b143ded3",  # stone
            "model_name": "stems",
            "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
            "metadata": {
                "multi_ids": [id + "_stem", id + "_complement"],
            },
        }
        model.stem.spawn(json.dumps(input))
    time.sleep(1000000)


def test_stems_v1():
    model = StemStub()

    # for an stem test
    id = str(uuid4())
    input = {
        "id": id,
        "prompt_audio": "a5e2198a-f352-4abb-9a24-7f81b143ded3",  # stone
        "model_name": "stems",
        "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "stem_type_id": 0,
            "multi_ids": [id + "_stem", id + "_complement"],
            # "stem_type_cfg_scale": 2,
        },
    }
    model.stem.remote(json.dumps(input))

    # test category
    id = str(uuid4())
    input = {
        "id": id,
        "prompt_audio": "a5e2198a-f352-4abb-9a24-7f81b143ded3",  # stone
        "model_name": "stems",
        "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "stem_type_group_name": "Vocals",
            "multi_ids": [id + "_stem", id + "_complement"],
        },
    }
    model.stem.remote(json.dumps(input))

    # test s3 url
    id = str(uuid4())
    input = {
        "id": id,
        "prompt_audio": "s3://suno-data-uploads/studio/uploads/a5e2198a-f352-4abb-9a24-7f81b143ded3.mp3",  # stone
        "model_name": "stems",
        "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "stem_type_id": 0,
            "multi_ids": [id + "_stem", id + "_complement"],
        },
    }
    model.stem.remote(json.dumps(input))


@app.local_entrypoint()
def main():
    if MODEL == "stems_v1":
        test_stems_v1()
    elif MODEL == "stems_v1_8_output":
        test_stems_v1_8_output()
    elif MODEL == "stems_v1_12_output":
        test_stems_v1_12_output()
        stress_test_stems_v1_12_output()
    else:
        raise ValueError(f"Unknown model: {MODEL}")
    print("Done")