"""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 enum import Enum
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, s3_object_exists
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, serialize_context
from ddtrace.propagation.http import HTTPPropagator
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_batched,
    encode_overlap,
    preload_models,
)
from suno_utils.utils.text import read_jsonl
from tqdm import tqdm
############## CHANGE THESE ##############

DEPLOYMENT_TYPE = "dev"  # dev, prod
MODEL_CONFIG = MODEL_CONFIG_DICT["stems_v1"]  # 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}"
TOKEN_QUEUE_NAME = f"token-queue-{DEPLOYMENT_TYPE}"
UPSAMPLE_DDOG_SERVICE = "stem-worker"

# 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)
apptoken_queue = modal.Queue.from_name(TOKEN_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()


TOKENS_PER_CHUNK = 250


class TokenSignalCode(Enum):
    STREAM_COMPLETE = -999.0
    GPT_ERROR = -998.0
    STEM_TIMED_OUT = -997.0


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=1000,  # 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=100 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()

        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:
                    t_queue_get_start = time.time()
                    item_id, audio_s3_path, history, parent_ctx_json = self.download_queue.get(
                        timeout=30
                    )
                    queue_wait_ms = int((time.time() - t_queue_get_start) * 1000)

                    try:
                        # activate parent trace context if present
                        try:
                            if parent_ctx_json:
                                ctx = HTTPPropagator.extract(json.loads(parent_ctx_json))
                                tracer.context_provider.activate(ctx)
                        except Exception:
                            pass

                        with tracer.trace("download_worker", service="stem-worker") as d_span:
                            d_span.set_tag("item_id", item_id)
                            d_span.set_tag("queue_wait_ms", queue_wait_ms)
                            d_span.set_tag("download_queue_size", self.download_queue.qsize())

                            if history is not None and history.history_latents is not None:
                                raw_audio = history.history_latents
                                duration = round(raw_audio.shape[0] / 25)
                                d_span.set_tag("source", "latent")
                                d_span.set_tag("duration_s", duration)
                                # Skip encoding for latents
                                t_ready_put = time.time()
                                self.ready_queue.put(
                                    (item_id, raw_audio, duration, True)
                                )  # is_latent=True
                                d_span.set_tag(
                                    "ready_queue_put_ms", int((time.time() - t_ready_put) * 1000)
                                )
                                d_span.set_tag("ready_queue_size", self.ready_queue.qsize())
                            else:
                                d_span.set_tag("source", "s3")
                                d_span.set_tag("audio_s3_path", audio_s3_path)
                                print(f"Downloading audio for {item_id}: {audio_s3_path}")
                                t_dl = time.time()
                                raw_audio = Audio.from_s3(audio_s3_path, n_channels=2)
                                d_span.set_tag("download_ms", int((time.time() - t_dl) * 1000))
                                # Safe audio props
                                try:
                                    d_span.set_tag("sr", getattr(raw_audio, "sample_rate", None))
                                    d_span.set_tag("n_channels", getattr(raw_audio, "n_channels", None))
                                except Exception:
                                    pass
                                duration = round(raw_audio.duration_s, 2)
                                d_span.set_tag("duration_s", duration)
                                # Submit for encoding
                                t_submit = time.time()
                                self.encode_executor.submit(
                                    self._encode_worker, item_id, raw_audio, duration, parent_ctx_json
                                )
                                d_span.set_tag("encode_submit_ms", int((time.time() - t_submit) * 1000))

                    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, parent_ctx_json: Optional[str] = None
    ):
        """Encodes raw audio (runs in encode_executor)."""
        try:
            # activate parent context if provided
            try:
                if parent_ctx_json:
                    ctx = HTTPPropagator.extract(json.loads(parent_ctx_json))
                    tracer.context_provider.activate(ctx)
            except Exception:
                pass

            with tracer.trace("encode_worker", service="stem-worker") as e_span:
                e_span.set_tag("item_id", item_id)
                e_span.set_tag("duration_s", duration)
                try:
                    e_span.set_tag("sr", getattr(raw_audio, "sample_rate", None))
                    e_span.set_tag("n_channels", getattr(raw_audio, "n_channels", None))
                    # Use array_float length if available
                    arr = getattr(raw_audio, "array_float", None)
                    if hasattr(arr, "shape"):
                        e_span.set_tag("samples", int(arr.shape[-1]))
                except Exception:
                    pass

                print(f"Encoding audio for {item_id}")
                t_enc = time.time()
                encoded_audio = encode_overlap(raw_audio, normalize_volume=False)
                e_span.set_tag("encode_ms", int((time.time() - t_enc) * 1000))
                try:
                    if hasattr(encoded_audio, "shape"):
                        e_span.set_tag("latents_shape", str(encoded_audio.shape))
                except Exception:
                    pass

                t_ready_put = time.time()
                self.ready_queue.put((item_id, encoded_audio, duration, False))  # is_latent=False
                e_span.set_tag("ready_queue_put_ms", int((time.time() - t_ready_put) * 1000))
                e_span.set_tag("ready_queue_size", self.ready_queue.qsize())
        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."""
        with tracer.trace("submit_audio_request", service="stem-worker") as s_span:
            s_span.set_tag("item_id", item.id)
            s_span.set_tag("has_history", history is not None)
            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://"):
                bucket = "suno-data-uploads"
                key_m4a = f"studio/uploads/{audio_s3_path}.m4a"
                key_opus = f"studio/uploads/{audio_s3_path}.opus"

                audio_s3_path = f"s3://suno-data-uploads/studio/uploads/{audio_s3_path}.opus"
                if s3_object_exists(bucket, key_m4a):
                    audio_s3_path = f"s3://{bucket}/{key_m4a}"
            s_span.set_tag("audio_s3_path", audio_s3_path)

            # Submit to download queue
            try:
                t_put = time.time()
                s_span.set_tag("download_queue_size_before", self.download_queue.qsize())
                parent_ctx_json = serialize_context()
                self.download_queue.put((item.id, audio_s3_path, history, parent_ctx_json), timeout=1)
                s_span.set_tag("download_queue_put_ms", int((time.time() - t_put) * 1000))
                s_span.set_tag("download_queue_size_after", self.download_queue.qsize())
            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."""
        with tracer.trace("get_ready_audio", service="stem-worker") as g_span:
            g_span.set_tag("item_id", item_id)
            g_span.set_tag("timeout_s", timeout)
            deadline = time.time() + timeout
            polls = 0
            empties = 0
            misses = 0
            requeues = 0
            t_wait_start = time.time()

            while time.time() < deadline:
                try:
                    polls += 1
                    with tracer.trace("get_ready_audio.queue_get", service="stem-worker") as q_span:
                        q_span.set_tag("attempt", polls)
                        q_span.set_tag("timeout_s", 1)
                        try:
                            q_span.set_tag("qsize_before", self.ready_queue.qsize())
                        except Exception:
                            pass
                        t_get_start = time.time()
                        ready_item_id, audio, duration, is_latent = self.ready_queue.get(timeout=1)
                        q_span.set_tag("get_wait_ms", int((time.time() - t_get_start) * 1000))
                        try:
                            q_span.set_tag("qsize_after", self.ready_queue.qsize())
                        except Exception:
                            pass

                    if ready_item_id == item_id:
                        # Clean up tracking
                        with tracer.trace("get_ready_audio.success_cleanup", service="stem-worker"):
                            with self.request_lock:
                                self.pending_requests.pop(item_id, None)
                        g_span.set_tag("wait_ms", int((time.time() - t_wait_start) * 1000))
                        g_span.set_tag("polls", polls)
                        g_span.set_tag("is_latent", bool(is_latent))
                        return audio, duration
                    else:
                        # Put back if it's for a different request
                        misses += 1
                        with tracer.trace(
                            "get_ready_audio.requeue_other", service="stem-worker"
                        ) as r_span:
                            r_span.set_tag("other_item_id", ready_item_id)
                            r_span.set_tag("miss_index", misses)
                            self.ready_queue.put((ready_item_id, audio, duration, is_latent))
                            try:
                                r_span.set_tag("ready_queue_size", self.ready_queue.qsize())
                            except Exception:
                                pass
                            requeues += 1

                except queue.Empty:
                    empties += 1
                    with tracer.trace("get_ready_audio.queue_empty", service="stem-worker") as e_span:
                        e_span.set_tag("attempt", polls)
                    continue

            g_span.set_tag("wait_ms", int((time.time() - t_wait_start) * 1000))
            g_span.set_tag("polls", polls)
            g_span.set_tag("empties", empties)
            g_span.set_tag("misses", misses)
            g_span.set_tag("requeues", requeues)
            raise TimeoutError(f"Audio not ready for {item_id} within {timeout}s")

    def code_generator(self, job: str):
        i = 0
        chunk_size = 125
        while True:
            # this is very performance sensitive. should investigate a faster method of ipc
            # codes = np.zeros((100, 14))
            codes = self.worker.engine_rpc_client.get_generated_codes(job, start=i)
            if codes is None:
                return
            i += len(codes)
            if isinstance(codes, np.ndarray):
                # Yield chunks of codes instead of individual codes
                for chunk_start in range(0, len(codes), chunk_size):
                    chunk_end = min(chunk_start + chunk_size, len(codes))
                    chunk = codes[chunk_start:chunk_end]
                    yield chunk  # Yields (250, 12, 128) or (remaining, 12, 128) for last chunk
            if i < TOKENS_PER_CHUNK:
                time.sleep(0.3)
            elif i < TOKENS_PER_CHUNK * 2:
                time.sleep(1)
            else:
                time.sleep(4)

    @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)
        with tracer.trace("audio_wait_phase", service="stem-worker"):
            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("gpu_diffusion_generation", service="stem-worker") as gen_span:
            pad_token = self.worker.model_config["cond_semantic_n_vocab"] - 1
            gen_span.set_tag("pad_token", pad_token)
            with tracer.trace(
                "gpu_diffusion_generation.add_request", service="stem-worker"
            ) as add_req_span:
                add_req_span.set_tag("tokens_shape", str((int(audio.shape[0]),)))
                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,
                    )
                )
                add_req_span.set_tag("job_id", job)
            with tracer.trace("gpu_diffusion_generation.wait_completion", service="stem-worker"):
                while not self.worker.engine_rpc_client.get_job_state(job)["completed"]:
                    time.sleep(0.2)
            with tracer.trace("gpu_diffusion_generation.get_results", service="stem-worker") as res_span:
                res_span.set_tag("job_id", job)
                latents = self.worker.engine_rpc_client.get_generated_codes(job)
                try:
                    res_span.set_tag("latents_shape", str(latents.shape))
                    res_span.set_tag("latents_nbytes", int(getattr(latents, "nbytes", 0)))
                except Exception:
                    pass
        print(f"Stem {item.id} got latents: {latents.shape}")

        item.metadata["audio_len"] = audio.shape[0]
        item.metadata["start_time"] = start_time

        # Capture parent context BEFORE dispatch span so decoder hangs off stem root, not this span
        parent_ctx_for_decoder = serialize_context()

        # Dispatch to decoder with tracing
        with tracer.trace("post_generation.dispatch", service="stem-worker") as disp_span:
            try:
                disp_span.set_tag("latents_shape", str(latents.shape))
                disp_span.set_tag("latents_nbytes", int(getattr(latents, "nbytes", 0)))
            except Exception:
                pass
            t_spawn = time.time()
            DecoderStub().generate.spawn(
                item.model_dump_json(), latents, parent_context=parent_ctx_for_decoder
            )
            disp_span.set_tag("spawn_ms", int((time.time() - t_spawn) * 1000))

        # Engine cleanup
        with tracer.trace("engine_job_cleanup", service="stem-worker"):
            self.worker.engine_rpc_client.remove_job(job)
        print(f"Done with stem code generation {item.id}")
        print(f"Stem part finished {item.id}: {(time.time() - start_time):.2f}s elapsed.")

    @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()


class DecoderWorker(S3Loader):
    def __init__(self):
        super().__init__()
        codec_path = diffusion_gen.get_model_if_needed(
            MODEL_CONFIG.codec_ckpt_path, cache_dir=MOUNT_PATH
        )
        preload_models(codec_path)


@app.cls(
    gpu=["A100-80GB"],  # H100 is faster, but if limited by resources, burst into A100
    cpu=12,
    secrets=SECRETS,
    timeout=1000,  # 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=200 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 DecoderStub:
    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 = DecoderWorker()
        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)

    @modal.method()
    @distributed_trace("decoder", UPSAMPLE_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE)
    def generate(self, item_json: str, latents: np.ndarray) -> None:
        item = QueueItem(**json.loads(item_json))
        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}",
        ]

        audio_len = item.metadata.get("audio_len")
        start_time = item.metadata.get("start_time")
        ids = item.multi_ids
        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 = []
        stride = 25 * 30
        if audio_len < stride:
            print(f"Stem {item.id} audio is too short, using stride {stride}")
            stride = 25

        # Decode all 12 stems in single GPU batch - pass the tensor directly
        print(f"Decoding {latents.shape[1]} stems in single GPU batch...")
        out_audios = decode_stream_to_full_audio_batched(latents, n_stride_tokens=stride)

        # Build outputs - same as before
        outputs = []
        for i, out_audio in enumerate(out_audios):
            if is_vocal_only:
                stem_id = f"{ids[0]}_stem_{i}"
            else:
                stem_id = ids[i]
            stem_queue_item = item.model_copy(update={"id": stem_id})
            latent = latents[:, i]
            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

        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=min(8, len(outputs))) 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
            # ]

            async def write_npz_async():
                """Write NPZ files asynchronously"""
                loop = asyncio.get_running_loop()
                npz_tasks = [
                    loop.run_in_executor(
                        executor,
                        self.worker._write_vae_latents_npz,
                        x_queue_item,
                        x,
                        MODEL_CONFIG.vae_version.value,
                        APP_NAME,
                    )
                    for x_queue_item, x, _ in outputs
                ]
                await asyncio.gather(*npz_tasks)

            async def write_all_parallel():
                """Write audio and NPZ files in parallel"""
                write_m4a = DEPLOYMENT_TYPE == "dev"
                await asyncio.gather(
                    # Audio writing
                    asyncio.gather(
                        *[
                            self.worker._write_audio_only_async(
                                stem_queue_item,
                                stem_audio,
                                s3_bucket="suno-data",
                                s3_folder="sara/persona_filter/minz_stems",
                                gain_adjust=MODEL_GAIN_ADJUST,
                                write_wav=False,
                                write_m4a=False,
                            )
                            for stem_queue_item, _, stem_audio in outputs_to_write_audio
                        ]
                    ),
                    # NPZ writing in parallel
                    # write_npz_async(),
                )

            # 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
            #         ],
            #     )

            audio_t0 = time.time()
            asyncio.run(write_all_parallel())
            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"]

        print(f"ids: {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 {}),
                },
            )

        # 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}")


def fill_local_token_queue_stem(local_token_queue: queue.Queue, partition: str):
    # Ensure the queue has not expired.
    if apptoken_queue.len(partition=partition) == 0:
        # Fail the decoding job for this partition.
        print(f"StemStub: No tokens in queue {partition}, putting GPT_ERROR")
        local_token_queue.put([TokenSignalCode.GPT_ERROR])
        return

    # Stream tokens into local queue.
    print(f"StemStub: Iterating over queue {partition}")
    for token_batch in apptoken_queue.iterate(partition=partition, item_poll_timeout=120):
        local_token_queue.put(token_batch)

        if (
            isinstance(token_batch[-1], TokenSignalCode)
            and token_batch[-1] == TokenSignalCode.STREAM_COMPLETE
        ):
            print(f"StemStub: Received STREAM_COMPLETE token for {partition}")
            return

    print(f"StemStub: No STREAM_COMPLETE token, putting STEM_TIMED_OUT for {partition}")
    local_token_queue.put([TokenSignalCode.STEM_TIMED_OUT])


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 = "bPUAcRELQgw"
    input = {
        "id": id,
        "prompt_audio": "s3://suno-data/datasets/harvest/genius_hq/audio/bPUAcRELQgw.webm",  # stone
        "model_name": "stems",
        # "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "stem_type_id": 0,
            "multi_ids": [id + "_vocals", id + "_complement"],
            # "stem_type_cfg_scale": 2,
        },
    }
    model.stem.remote(json.dumps(input))

    # test category
    id = "xoADLQ1CwUA"
    input = {
        "id": id,
        "prompt_audio": "s3://suno-data/datasets/harvest/genius_hq/audio/xoADLQ1CwUA.webm",  # stone
        "model_name": "stems",
        # "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "stem_type_group_name": "Vocals",
            "multi_ids": [id + "_vocals", id + "_complement"],
        },
    }
    model.stem.remote(json.dumps(input))

    # test s3 url
    id = "7SbMtVQmThY"
    input = {
        "id": id,
        "prompt_audio": "s3://suno-data/datasets/harvest/genius_hq/audio/7SbMtVQmThY.webm",  # stone
        "model_name": "stems",
        # "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/",
        "metadata": {
            "stem_type_id": 0,
            "multi_ids": [id + "_vocals", id + "_complement"],
        },
    }
    model.stem.remote(json.dumps(input))


def batch_process_stems():
    print("Batch process with multiprocessing")
    model = StemStub()

    stem_data = read_jsonl("/home/sara/sara/leftover_vox.jsonl")
    print(f"Processing {len(stem_data)} stems...")

    for stem in tqdm(stem_data):
        try:
            s3_path = stem["s3_filepath"]
            is_valid = stem.get("valid_artist_vox", False)
            vox_filepath = stem.get("vox_filepath", None)

            if is_valid and stem["duration_s"] < 300:
                song_id = stem["id"]
                # s3_check_path = f"s3://suno-data/sara/persona_filter/minz_stems/{song_id}_vocals.opus"
                # s3_check_path_2 = f"s3://suno-data/sara/persona_filter/minz_stems/{song_id}_vocal.opus"
                if True:  # not check_s3_file_exists(s3_check_path) and not check_s3_file_exists(s3_check_path_2):
                    # if True:
                    input_data = {
                        "id": song_id,
                        "prompt_audio": s3_path,
                        "metadata": {
                            "stem_type_group_name": "Vocals",
                            "multi_ids": [song_id + "_vocals", song_id + "_complement"],
                        },
                    }
                    model.stem.spawn(json.dumps(input_data))
                    # time.sleep(0.01)
                    # print(f"Processed: {song_id}")
                else:
                    print(f"Skipped (exists): {song_id}")
            else:
                print(f"Skipped (invalid): {stem['id']}")

        except Exception as e:
            print(f"Error processing {stem.get('id', 'unknown')}: {str(e)}")

    time.sleep(1000000)


@app.local_entrypoint()
def main():
    if MODEL == "stems_v1":
        # test_stems_v1()
        batch_process_stems()
    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}")

    time.sleep(60 * 2)
    print("Done")