import os

import time
import modal
import torch
import pathlib
import funcy
import pandas as pd
import numpy as np
import tempfile

from suno_utils.gpt.generation import GenerationConfig
from suno_utils.gpt.engine import Engine
from suno_utils.gpt.generation_engine import make_request
from suno_utils.worker.settings import s3_client
from suno_utils.worker.modal_base import MODAL_MOUNTS
from suno_utils.diffusion import generation as diffusion_gen
from suno_utils.gpt import chirp_v2_5 as chirp_v2

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",
        ]
    )
    .dockerfile_commands(
        [
            "COPY --from=datadog/serverless-init:1.2.1 /datadog-init /app/datadog-init",
            'ENTRYPOINT ["/app/datadog-init"]',
        ]
    )
    .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",
        "funcy",
        "fasttext",
    )
    .pip_install_from_pyproject(
        str(
            "/home/christian/code/glockenspiel/suno_utils/pyproject.toml",
        ),
        # force_build=True,
    )
    .pip_install("torch==2.4.0")  # this is cause flash-attn can't work with 2.5 yet
    .run_commands(
        "FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn==2.6.3 --no-build-isolation",
    )
    .pip_install(
        "torch==2.5.1",
        "torchaudio==2.5.1",
        "transformers==4.30.0",
        "numpy==1.26.4",
    )
)


class GenerateWorker:
    def __init__(self, gpt_ckpt: str, dit_model_filepath: str, output_path: str):
        self.output_path = output_path
        start_time = time.time()
        print("Start loading models")

        gpt_ckpt_path = chirp_v2._get_model_if_needed(gpt_ckpt, cache_dir=MOUNT_PATH)
        tokenizer_path = chirp_v2._get_model_if_needed(
            chirp_v2.TOKENIZER_PATH, cache_dir=MOUNT_PATH
        )
        print(tokenizer_path)

        N_BATCH = 2
        engine = Engine(
            gpt_ckpt_path,
            tokenizer_path=tokenizer_path,
            max_sequences=4 * N_BATCH,
            compile=False,
        )
        cfg = engine.model.config

        self.engine = engine
        self.cfg = cfg

        # load gpt model
        num_gpus = torch.cuda.device_count()
        cuda_device = torch.cuda.current_device()
        print(f"Found {num_gpus} GPUs. Using GPU {cuda_device}.")

        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(
            dit_model_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,
            compile=True,  # can't compile with flash v2)
        )

        # diff_model_fp = "/app/suno/checkpoints/2025-01-19_05-10-13_s1983/step_3000_infer.pt"  # best
        # diff_model_fp = "/app/suno/checkpoints/2024-12-16_15-54-26_s3031/last_ckpt_infer.pt"
        # diff_model_fp = "s3://suno-data/georg/tmp/2b_prefix_ft.pt"

        # codec_filepath = "s3://suno-data/minz/models/dac_vae_fixed_25hz.pth"
        # codec_filepath = "s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth"
        # tokenizer_filepath = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json"
        # semantic_model_filepath = "s3://suno-data/georg/models/semantic/mert_25.pt"
        # semantic_clusters_filepath = (
        #    "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy"
        # )

        if "dac_vae_fixed" in codec_path or "dac_vae_tuned" in codec_path:
            patch_size = 1
            codec_scale_factor = 0.4
        elif "convnext_vae_tuned" in codec_path:
            patch_size = 1
            codec_scale_factor = 1.0
        else:
            patch_size = 1
            codec_scale_factor = 2.5

        print(
            f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds"
        )

    @staticmethod
    def download_models(gpt_ckpt, dit_model_filepath, 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(dit_model_filepath, cache_dir=dir_path)

        _ = chirp_v2._get_model_if_needed(gpt_ckpt, cache_dir=dir_path)
        _ = chirp_v2._get_model_if_needed(chirp_v2.TOKENIZER_PATH, cache_dir=dir_path)

        print("Finish downloading models")

    def generate(self, work_item):
        for item in work_item:

            n_skip_semantic = 1

            gconf = GenerationConfig(
                text=item["text"],
                text_tags=item["tags"],
                cfg_coef=1.0,
                cfg_coef_tags=1.0,
                cfg_coef_max_steps=None,
                cfg_coef_tags_max_steps=None,
                n_repeat_tags=1,
                n_skip_semantic=1,
                # text_start_control_tags="{duration:210}",
                temp_semantic=0.92,
                n_batch=1,
                min_text_offset=0,
                eos_pad_duration_s=0,
                max_gen_duration_s=int(8 * 60 / n_skip_semantic),
                random_seed=np.random.randint(0, 1000000),
            )

            requests = [
                make_request(
                    f"{i}", gconf, self.engine.model.config, self.engine.tokenizer
                )
                for i in range(1)
            ]
            jobs = self.engine.run_request(requests, tqdm_enabled=True)
            out_gpt = []
            for n, job in enumerate(jobs):
                stream = self.engine.token_generator(job)
                arr = torch.stack(list(stream))[:, 1]
                if arr[-1] == 4000:
                    arr = arr[:-1]
                print(f"{round(arr.shape[-1]/25*n_skip_semantic)}s for track {n}")
                # do stuff incase skip
                arr2 = (
                    torch.zeros(arr.shape[0] * n_skip_semantic, dtype=arr.dtype)
                    + self.cfg.semantic_pad_token
                )
                arr2[::n_skip_semantic] = arr
                # add
                out_gpt.append(arr2)

            # perform upsampling with multiple seeds
            diffusion_inference_config = diffusion_gen.DiffusionGenerationConfig(
                audio=out_gpt[0],
                lyrics=item["text"],
                tags=item["tags"],
                text_cfg_coef=1.5,
                ctx_cfg_coef=1.0,
                steps=16,
                seed=np.random.randint(0, 1000000),
            )

            upsampled_audio = diffusion_gen.generate(diffusion_inference_config)

            # we want to save out
            # 3. mp3 audio file

            with tempfile.TemporaryDirectory() as td:
                base_filename = f"{item['id']}"

                # save audio to s3
                upsampled_audio_path = os.path.join(td, f"{base_filename}.mp3")
                upsampled_audio.write_hq_mp3(upsampled_audio_path)

                s3_filepath = os.path.join(self.output_path, f"{base_filename}.mp3")
                print(f"Uploading to {s3_filepath}")
                s3_client.upload_file(
                    upsampled_audio_path,
                    "suno-data",
                    s3_filepath,
                    ExtraArgs={
                        "ContentType": "audio/mp3",
                    },
                )


# base model 6b
GPT_CKPT = (
    "s3://suno-data/christian/checkpoints/gpt/2025-02-04_21-04-31-last_ckpt_infer.pt"
)

# 6b v2 finetune
# GPT_CKPT = (
#    "s3://suno-data/christian/checkpoints/gpt/2025-02-10_16-52-41-step_9000_infer.pt"
# )

# GPT_CKPT = "s3://suno-data/christian/checkpoints/gpt/2025-02-13_11-27-29-step_10_000.pt"

# prod diffusion model
DIT_MODEL_FILEPATH = "s3://suno-data/tony/tmp/diff/dit_v3_dpo_t10_3k_5e6_b100_t25.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/georg/tmp/2b_prefix_ft.pt"


def download_model_wrapper_d():
    # this print is necessary to have modal rerun this when MODEL changes
    # Modal tracks referenced global variables
    # Change the name of the function to force a rerun
    print("Downloading model for history encoder")
    GenerateWorker.download_models(GPT_CKPT, DIT_MODEL_FILEPATH)


image = base_image.run_function(download_model_wrapper_d, secrets=SECRETS)
APP_NAME = f"batch-generate-eval-gpu"
app = modal.App(APP_NAME, image=image, secrets=SECRETS)

N_MAX_REPLICAS = 8


@app.cls(
    gpu=modal.gpu.H100(count=1),
    cpu=4,
    secrets=SECRETS,
    timeout=2 * 60 * 60,
    container_idle_timeout=240,
    mounts=MODAL_MOUNTS,
    memory=15000,
    concurrency_limit=N_MAX_REPLICAS,
)
class GenerateStub:
    def __init__(self, gpt_ckpt: str, dit_ckpt: str, output_path: str):
        import torch

        num_gpus = torch.cuda.device_count()
        print(f"Found {num_gpus} GPUs.")
        self.worker = GenerateWorker(gpt_ckpt, dit_ckpt, output_path)

    @modal.method()
    def generate(self, work_item: list[dict]):
        return self.worker.generate(work_item)


@app.local_entrypoint()
def main():
    # checkpoints in s3://suno-data/christian/checkpoints
    # prompts in s3://suno-data/christian/prompts
    # outputs in s3://suno-data/christian/outputs

    ckpt_name = GPT_CKPT.split("/")[-1].replace(".pt", "")
    print(ckpt_name)

    # load the prompts to evaluate
    prompts_path = "s3://suno-data/christian/prompts/suno_bench_prompts_v1.csv"
    # prompts_path = "s3://suno-data/christian/prompts/02_singing_prompts.csv"

    # load the prompts to evaluate
    prompts = pd.read_csv(prompts_path)
    print(len(prompts))

    chunksize = 16  # num of prompts per worker
    work_items = list(funcy.chunks(chunksize, prompts.to_dict(orient="records")))
    print(len(work_items), "work items")

    worker = GenerateStub(
        GPT_CKPT, DIT_MODEL_FILEPATH, f"christian/outputs/{ckpt_name}"
    )

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

    print("Running batch inference...")
    t0 = time.time()
    _ = list(worker.generate.map(work_items))
    print(round((time.time() - t0) / 60 / 60), "h total for batch generation")