import funcy
import gc
import os
import pathlib
import re
import tempfile
import time

import modal
from modal import enter
import numpy as np

from suno_utils.utils.display import suppress_logging
from contextlib import redirect_stderr
from suno_utils.utils.s3 import read_from_s3, upload_s3_files, list_s3_dir
from suno_utils.utils.text import read_jsonl, write_jsonl
from suno_utils.tasks.data_loader import load_audio_mp

from suno_utils.diffusion import generation as diffusion_gen
from suno_utils.worker.settings import s3_client

N_MAX_REPLICAS = 1
MOUNT_PATH = "/suno/models"

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("api-callback-token"),
    modal.Secret.from_name("datadog-metrics"),
]

base_image = (
    modal.Image.debian_slim()
    .apt_install(
        "curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3", "zlib1g-dev", "git"
    )
    .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.4.0")
    .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",
        "rpyc",
        "biopython>=1.81",  # TODO: don't love this depdendency, for hoot
        "pynvml",  # for torch cuda utilization
        "torchsde",
    )
    .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==2.5.9.post1 --no-build-isolation"
    )
    .pip_install(
        "torch==2.5.0.dev20240728+cu121",
        "torchaudio==2.4.0.dev20240728+cu121",
        index_url="https://download.pytorch.org/whl/nightly/cu121",
    )
)

stub = modal.App("batch-upsample-gpu", image=base_image)


# ## Defining the prediction function
#
# * Container lifecycle hook: this lets us load the model only once in each container
#
# Instead of a using `@stub.function()` in the global scope,
# we put the method on a class, and define an `__enter__` method on that class.
# Modal reuses containers for successive calls to the same function, so
# we want to take advantage of this and avoid setting up the same model
# for every function call.

# Notes:
# relative speeds: 5.5s on a5000, 7.5s on A10G, 18.5s on T4
# n_cpu on T4/A10G: 24
# audio load: 75s for 1k files on 24 cores (70h audio) ~3-4k realtime


@stub.cls(
    gpu="A10G",
    timeout=2 * 60 * 60,
    concurrency_limit=N_MAX_REPLICAS,
    secrets=[modal.Secret.from_name("aws-bucket")],
    # cloud="oci",
    memory=150000,  # youtube 5k needs about 100gb
    retries=3,
)
class Worker:
    def __init__(
        self,
        encode_s3_dir,
        min_duration_s,
        max_duration_s,
    ):
        """
        Args:
            encode_s3_dir: S3 directory to store upsampled audio files
            min_duration_s: Minimum duration of audio to consider
            max_duration_s: Maximum duration of audio to consider
        """
        self.encode_s3_dir = encode_s3_dir
        self.min_duration_s = min_duration_s
        self.max_duration_s = max_duration_s

    @enter()
    def set_up_models(self):
        start_time = time.time()
        print("Start loading models")
        tokenizer_path = diffusion_gen.get_model_if_needed(
            diffusion_gen.TOKENIZER_FILEPATH, cache_dir=MOUNT_PATH
        )
        semantic_model_path = diffusion_gen.get_model_if_needed(
            diffusion_gen.SEMANTIC_MODEL_FILEPATH, cache_dir=MOUNT_PATH
        )
        semantic_clusters_path = diffusion_gen.get_model_if_needed(
            diffusion_gen.SEMANTIC_CLUSTERS_FILEPATH, cache_dir=MOUNT_PATH
        )
        codec_path = diffusion_gen.get_model_if_needed(
            diffusion_gen.CODEC_FILEPATH, cache_dir=MOUNT_PATH
        )
        dit_model_path = diffusion_gen.get_model_if_needed(
            diffusion_gen.DIT_MODEL_FILEPATH, cache_dir=MOUNT_PATH
        )
        dit_config_path = diffusion_gen.get_model_if_needed(
            diffusion_gen.DIT_CONFIG_FILEPATH, cache_dir=MOUNT_PATH
        )
        diffusion_gen.preload_models(
            tokenizer_filepath=tokenizer_path,
            semantic_model_filepath=semantic_model_path,
            semantic_clusters_filepath=semantic_clusters_path,
            codec_filepath=codec_path,
            dit_model_filepath=dit_model_path,
            dit_config_filepath=dit_config_path,
            compile=True,  # can't compile with flash v2)
        )
        print(
            f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds"
        )

    @staticmethod
    def download_models(dir_path=MOUNT_PATH):
        """Download diffusion models."""
        print("Start downloading models")
        _ = diffusion_gen.get_model_if_needed(
            diffusion_gen.TOKENIZER_FILEPATH, cache_dir=dir_path
        )
        _ = diffusion_gen.get_model_if_needed(
            diffusion_gen.SEMANTIC_MODEL_FILEPATH, cache_dir=dir_path
        )
        _ = diffusion_gen.get_model_if_needed(
            diffusion_gen.SEMANTIC_CLUSTERS_FILEPATH, cache_dir=dir_path
        )
        _ = diffusion_gen.get_model_if_needed(
            diffusion_gen.CODEC_FILEPATH, cache_dir=dir_path
        )
        _ = diffusion_gen.get_model_if_needed(
            diffusion_gen.DIT_MODEL_FILEPATH, cache_dir=dir_path
        )
        _ = diffusion_gen.get_model_if_needed(
            diffusion_gen.DIT_CONFIG_FILEPATH, cache_dir=dir_path
        )
        print("Finish downloading models")

    @modal.method()
    def upsample(self, work_item):
        pass


if __name__ == "__main__":

    encode_s3_dir = os.path.join(base_s3_dir, output_name)

    print("Downloading data...")
    metas = get_metas(os.path.join(base_s3_dir, "metas.jsonl"))
    print(len(metas), "data items loaded")

    work_items = list(funcy.chunks(chunksize, metas))
    work_items = list(zip(range(len(work_items)), work_items))
    print(len(work_items), "work items")

    worker = Worker(
        encode_s3_dir,
        min_duration_s,
        max_duration_s,
    )

    print("Testing inference...")
    t0 = time.time()
    for work_item in work_items[:1]:
        _ = worker.embed.remote(work_item)
    print(f"{int(round(time.time()-t0))}s for test")

    print("Running batch inference...")
    t0 = time.time()
    _ = list(worker.embed.map(work_items[1:]))
    print(round((time.time() - t0) / 60 / 60), "h total for batch embed")