import os

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

from suno_utils.audio import Audio
from suno_utils.utils.s3 import read_from_s3
from suno_utils.utils.text import read_jsonl
from suno_utils.worker.settings import s3_client
from suno_utils.worker.modal_base import MODAL_MOUNTS
from suno_utils.utils.s3 import read_from_s3, upload_s3_files, list_s3_dir
from suno_utils.diffusion.generation import (
    generate_offline,
    preload_models,
    get_model_if_needed,
    TOKENIZER_FILEPATH,
    SEMANTIC_MODEL_FILEPATH,
    SEMANTIC_CLUSTERS_FILEPATH,
    _retrieve_models,
)

from suno_utils.tasks.mert_25 import (
    preload_models as preload_semantic_models,
    encode as encode_semantic,
    EMBEDDING_RATE as SEMANTIC_HZ,
)

from suno_utils.tasks.dac_vae_fixed_25hz import (
    preload_models as preload_dac_vae_fixed_25hz,
    encode as codec_encode,
    decode as codec_decode,
)

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(
        "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",
        "wandb",
    )
    .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.44.0",
        "numpy==1.26.4",
    )
)


class GenerateWorker:
    def __init__(
        self,
        dit_model_filepath: str,
        codec_model_filepath: 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}.")

        if (
            "dac_vae_fixed" in codec_model_filepath
            or "dac_vae_tuned" in codec_model_filepath
        ):
            self.patch_size = 1
            self.codec_scale_factor = 0.4
            self.downscale_ctx_vector = False
        elif "convnext_vae_tuned" in codec_model_filepath:
            self.patch_size = 1
            self.codec_scale_factor = 1.0
            self.downscale_ctx_vector = False
        else:
            self.patch_size = 1
            self.codec_scale_factor = 2.5
            self.downscale_ctx_vector = True

        _ = preload_models(
            tokenizer_filepath=TOKENIZER_FILEPATH,
            semantic_model_filepath=SEMANTIC_MODEL_FILEPATH,
            semantic_clusters_filepath=SEMANTIC_CLUSTERS_FILEPATH,
            codec_filepath=codec_model_filepath,
            dit_model_filepath=dit_model_filepath,
            weights_precision=torch.bfloat16,
            model_type="prefix",
            codec_scale_factor=self.codec_scale_factor,
        )

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

    @staticmethod
    def download_models(dit_model_filepath, codec_model_filepath, dir_path=MOUNT_PATH):
        """Download diffusion models."""
        print("Start downloading models")
        _ = get_model_if_needed(TOKENIZER_FILEPATH, cache_dir=dir_path)
        _ = get_model_if_needed(SEMANTIC_MODEL_FILEPATH, cache_dir=dir_path)
        _ = get_model_if_needed(SEMANTIC_CLUSTERS_FILEPATH, cache_dir=dir_path)
        _ = get_model_if_needed(codec_model_filepath, cache_dir=dir_path)
        _ = get_model_if_needed(dit_model_filepath, cache_dir=dir_path)

        print("Finish downloading models")

    def generate(self, work_items):
        print(f"Generating {len(work_items)} work items")

        # predfine 100 seeds
        seeds = np.random.randint(0, 1000000, size=100)

        for work_item in work_items:
            print(work_item)
            # get the id
            gen_id = work_item["id"]
            # check if the metadata exists on s3

            # skip anything if the duration is longer than 240s
            # if work_item["duration_s"] > 480:
            #    print(f"Skipping {gen_id} because it is longer than 480s")
            #    continue

            # check if the metadata.json file exists on s3
            # metadata_path = os.path.join(self.output_path, f"{gen_id}_metadata.json")

            try:
                s3_filepath = f"s3://suno-data-uploads/studio/uploads/{gen_id}.npz"
                data = read_from_s3(s3_filepath, read_f=np.load)

                if "v3.0_raw" in data:
                    codes = data["v3.0_raw"]
                elif "v3.5_raw" in data:
                    codes = data["v3.5_raw"]
                elif "v4.0_raw" in data:
                    codes = data["v4.0_raw"]
                else:
                    raise ValueError("No codes found")

                semantic_codes = torch.from_numpy(codes[:, 0]).long()  # .cuda()
            except Exception as e:
                print(f"Error loading semantic codes: {e}")
                continue

            # perform upsampling with multiple seeds

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

            sampled_seeds = np.random.choice(seeds, size=10)

            try:
                with tempfile.TemporaryDirectory() as td:
                    base_filename = f"{work_item['id']}"
                    for seed in sampled_seeds:
                        # create a uuid for this example
                        upsample_uuid = str(uuid.uuid4())
                        seed = int(seed)
                        print(f"Upsampling with seed {seed}")

                        steps = np.random.choice([6, 8, 10, 12, 14, 16])
                        text_cfg_coef = np.random.choice(
                            [1.0, 1.25, 1.5, 2.0, 2.5, 3.0]
                        )
                        noise_ctx_vector = np.random.choice([0.0, 0.1, 0.2, 0.5])

                        # check for valid inputs
                        if work_item["text"] is None:
                            text = ""
                        else:
                            text = work_item["text"]

                        if work_item["tags"] is None:
                            tags = ""
                        else:
                            tags = work_item["tags"]

                        # generate normally using ctx
                        upsampled_audio, upsampled_latents = generate_offline(
                            audio=semantic_codes,
                            lyrics=text,
                            tags=tags,
                            text_cfg_coef=text_cfg_coef,
                            steps=steps,
                            sigma_max=50.0,
                            downscale_ctx_vector=self.downscale_ctx_vector,
                            noise_ctx_vector=noise_ctx_vector,
                            sampler_eta=1.0,
                            sampler_s_noise=1.0,
                            return_latents=True,
                            seed=seed,
                            use_ctx_vector=True,
                        )

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

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

                        # save vae latents to s3
                        upsampled_latents_path = os.path.join(
                            td, f"{base_filename}_upsample_{upsample_uuid}_vae.npz"
                        )
                        np.savez(upsampled_latents_path, upsampled_latents)

                        s3_filepath = os.path.join(
                            self.output_path,
                            f"{base_filename}_upsample_{upsample_uuid}_vae.npz",
                        )
                        print(f"Uploading to {s3_filepath}")
                        s3_client.upload_file(
                            upsampled_latents_path,
                            "suno-data",
                            s3_filepath,
                        )

                        # save semantic codes to s3
                        semantic_codes_path = os.path.join(
                            td, f"{base_filename}_upsample_{upsample_uuid}_semantic.npz"
                        )
                        np.savez(semantic_codes_path, semantic_codes)

                        s3_filepath = os.path.join(
                            self.output_path,
                            f"{base_filename}_upsample_{upsample_uuid}_semantic.npz",
                        )
                        print(f"Uploading to {s3_filepath}")
                        s3_client.upload_file(
                            semantic_codes_path,
                            "suno-data",
                            s3_filepath,
                        )

                        # create a json file with the metadata
                        metadata = {
                            "clip_id": work_item["id"],
                            "seed": seed,
                            "text_cfg_coef": float(text_cfg_coef),
                            "text": text,
                            "tags": tags,
                            "steps": int(steps),
                            "noise_ctx_vector": float(noise_ctx_vector),
                        }
                        metadata_path = os.path.join(td, f"{base_filename}.json")
                        with open(metadata_path, "w") as f:
                            json.dump(metadata, f)
                        s3_filepath = os.path.join(
                            self.output_path,
                            f"{base_filename}_upsample_{upsample_uuid}_metadata.json",
                        )
                        print(f"Uploading to {s3_filepath}")
                        s3_client.upload_file(
                            metadata_path,
                            "suno-data",
                            s3_filepath,
                            ExtraArgs={
                                "ContentType": "application/json",
                            },
                        )
            except Exception as e:
                print(f"Error generating: {e}")
                return


# prod diffusion model
# DIT_MODEL_FILEPATH = "s3://suno-data/georg/tmp/2b_prefix_ft.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/tony/tmp/diff/dit_v3_dpo_t10_3k_5e6_b100_t25.pt"
DIT_MODEL_FILEPATH = (
    "s3://suno-data/christian/checkpoints/diffusion/v45_2b_step_2_600_000.pt"
)
CODEC_MODEL_FILEPATH = "s3://suno-data/minz/models/dac_vae_tuned_25hz.pth"


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(DIT_MODEL_FILEPATH, CODEC_MODEL_FILEPATH)


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

N_MAX_REPLICAS = 8


@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, dit_ckpt: str, codec_ckpt: str, output_path: str):
        import torch

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

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


@app.local_entrypoint()
def main():

    ckpt_name = os.path.basename(DIT_MODEL_FILEPATH).split(".")[0]
    print(ckpt_name)

    # load the positive prompts
    work_items = read_from_s3(
        "s3://suno-data/christian/sft/pos_interesting_clips_up_u_1_20241201_full.jsonl",
        read_f=read_jsonl,
    )

    # just do the first 100
    work_items = work_items[:128]

    # work_items = read_jsonl(
    #    "/home/christian/code/christian/metadata/discogs_subset_sampled_metas.jsonl"
    # )

    chunksize = 16  # num of prompts per worker
    print(len(work_items), "work items")

    worker = GenerateStub(
        DIT_MODEL_FILEPATH,
        CODEC_MODEL_FILEPATH,
        f"christian/outputs/{ckpt_name}/10_apr_2025",
    )

    # look for existing metadata files
    if False:
        existing_files = list_s3_dir(
            f"s3://suno-data/christian/outputs/{ckpt_name}/pos_interesting_clips_up_u_1_20241201_full"
        )
        existing_files = [
            f[0] for f in existing_files if f[0].endswith("_metadata.json")
        ]
        existing_ids = [
            os.path.basename(f).split("_metadata.json")[0] for f in existing_files
        ]
        existing_ids = set(existing_ids)

        # now filter the work items
        work_items = [w for w in work_items if w["id"] not in existing_ids]

    # split into chunks
    work_items = [
        work_items[i : i + chunksize] for i in range(0, len(work_items), chunksize)
    ]

    print(f"{len(work_items)} chunks of {chunksize} work items")

    if False:
        print("Testing inference...")
        t0 = time.time()
        for work_item in work_items[:1]:
            _ = worker.generate.remote(work_item)
        print(f"{int(round(time.time()-t0))}s for test")
        # Calculate predicted time based on number of chunks and available GPUs
        num_chunks = len(work_items) - 1  # Excluding the test chunk

        time_per_chunk = time.time() - t0
        predicted_time = (
            (time_per_chunk * num_chunks)
            / (N_MAX_REPLICAS if N_MAX_REPLICAS > 0 else 1)
            / 60
            / 60
        )
        print(f"predicted time: {predicted_time}h for batch generation")

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


# aws s3 ls s3://suno-data/christian/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas/
# aws s3 sync s3://suno-data/christian/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas/