import funcy
import gc
import os
import json
import pathlib
import re
import tempfile
import time
import pandas as pd

import torch
import modal
from modal import enter
import numpy as np
import torchaudio
from itertools import combinations

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.gpt.generation import GenerationConfig
from suno_utils.gpt.engine import Engine
from suno_utils.gpt.generation_engine import make_request
from suno_utils.diffusion.generation import preload_diff_models, generate

from suno_utils.worker.modal_base import MODAL_MOUNTS


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",
    )
    .pip_install_from_pyproject(
        str(pathlib.Path(__file__).parent.parent.parent.parent / "pyproject.toml"),
    )
    .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",
    )
)


def setup():

    N_BATCH = 2
    engine = Engine(
        "/app/suno/models/chirp_v2/tokenizer_60k.json",
        max_sequences=4 * N_BATCH,
        compile=False,
    )
    cfg = engine.model.config

    diff_model_fp = (
        "/app/suno/checkpoints/2024-12-10_19-53-58_s7792/step_3000_infer.pt"  # prod
    )
    # 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"

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

    _ = preload_diff_models(
        tokenizer_filepath="/home/georg/notebooks/gpu_nb/tmp/tokenizer_60k.json",
        semantic_model_filepath="/home/georg/notebooks/gpu_nb/tmp/mert_25.pt",
        semantic_clusters_filepath="/home/georg/notebooks/gpu_nb/tmp/mert_25_2x4k.npy",
        codec_filepath=codec_filepath,
        dit_model_filepath=diff_model_fp,
        model_type="prefix",
        weights_precision=torch.bfloat16,
        compile=False,
        # codec_scale_factor=codec_scale_factor,
    )


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

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

        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
        )

        min_batch_size = min(32, max_sequences)
        if MODEL_CONFIG.min_batch_size is not None:
            min_batch_size = MODEL_CONFIG.min_batch_size

        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
        )

        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=False,  # 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
        )
        print("Finish downloading models")

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

            # perform upsampling with multiple seeds
            for seed in seeds:
                diffusion_inference_config = diffusion_gen.DiffusionGenerationConfig(
                    audio=semantic_codes,
                    lyrics=item["text"],
                    tags=item["tags"],
                    text_cfg_coef=4.0,
                    ctx_cfg_coef=1.0,
                    steps=steps,
                    seed=seed,
                )

                try:
                    upsampled_latents, upsampled_audio = diffusion_gen.offline_generate(
                        semantic_codes, diffusion_inference_config
                    )
                    upsampled_audios.append(upsampled_audio)
                except Exception as e:
                    print(f"Error generating {s3_filepath}: {e}. Skipping...")
                    continue

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

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

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

                    s3_client.upload_file(
                        upsampled_audio_path,
                        self.s3_bucket,
                        f"{self.encode_s3_dir}/{item_id}/{base_filename}.mp3",
                        ExtraArgs={
                            "ContentType": "audio/mp3",
                        },
                    )


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


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

N_MAX_REPLICAS = 64


@app.cls(
    gpu=modal.gpu.A10G(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, diff_ckpt: str, output_path: str):
        import torch

        num_gpus = torch.cuda.device_count()
        print(f"Found {num_gpus} GPUs.")
        self.worker = GenerateWorker(gpt_ckpt, diff_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

    # gpt & diffusion checkpoints to evaluate, need to move ckpts to s3 first
    gpt_ckpt = (
        "s3://suno-data/christian/checkpoints/025-02-02_21-04-52-last_ckpt_infer.pt"
    )
    diff_ckpt = "s3://suno-data/tony/tmp/diff/dit_v3_dpo_t10_3k_5e6_b100_t25.pt"

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

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

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

    worker = GenerateStub(gpt_ckpt, diff_ckpt, "s3://suno-data/christian/outputs")

    # print("Testing inference...")
    # t0 = time.time()
    # for work_item in work_items[:1]:
    #     _ = worker.upsample.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 embed")