import json
import time
from uuid import uuid4

import modal

from suno_utils.worker.chirp_worker_v1 import ChirpV1Worker
from suno_utils.worker.modal_base import MODAL_MOUNTS, get_modal_base_image
from suno_utils.worker.utils import recursive_ls_dir

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/",
        }
    ),
]


def download_model_wrapper_d5():
    ChirpV1Worker.download_models()


def download_whisper():
    import whisper

    whisper.load_model("small.en")
    whisper.load_model("small")

    recursive_ls_dir("/suno/models")


image = (
    get_modal_base_image().run_function(download_model_wrapper_d5, secrets=SECRETS)
    # .run_function(download_whisper, secrets=SECRETS)
)
STUB_NAME = "chirp-v1-alpha"
stub = modal.Stub(STUB_NAME, image=image)


@stub.cls(
    cpu=2.0,
    # memory=16384,
    gpu=modal.gpu.A10G(count=1),
    secrets=SECRETS,
    timeout=800,
    container_idle_timeout=400,
    mounts=MODAL_MOUNTS,
    retries=modal.Retries(
        max_retries=1,
        backoff_coefficient=2.0,
        initial_delay=10.0,
    ),
    keep_warm=1,
    concurrency_limit=120,
)
class ChirpV1Stub:
    def __enter__(self):
        import torch

        num_gpus = torch.cuda.device_count()
        print(f"Found {num_gpus} GPUs.")
        recursive_ls_dir("/suno/models")

        self.worker = ChirpV1Worker(0, bg_image="/suno/models/assets/wave-bg-2.png")
        self.worker.preload()

    @modal.method()
    def generate(self, queue_item: str):
        import json

        from suno_utils.worker.schema import QueueItem

        print(queue_item)
        start_time = time.time()
        item = QueueItem(**json.loads(queue_item))
        item_id = item.id
        print(item.metadata)
        is_square = item.metadata["is_square"] if "is_square" in item.metadata.keys() else False
        ids = []

        f = modal.Function.lookup("sdxl-prod", "StableDiffusion.generate_image_item")
        fn_call = f.spawn(item.copy(deep=True, update={"callback_url": None}).json())

        try:
            audios, raw_arrays = self.worker.process_item(item)
            items = []
            for i, audio in enumerate(audios):
                new_id = f"{item_id}_{i}"
                ids.append(new_id)
                item.id = new_id
                # item.prompt_text = trimmed[i]
                items.append(item.json())
                self.worker._write_audio_only(
                    item,
                    audio,
                )

                self.worker._write_npz(item, raw_arrays[i], "chirp-v1", "1.0.0.0")

            image_url = fn_call.get(timeout=None)
            print("image located", f"image_{item_id}.png")

            video_time = time.time()
            f = modal.Function.lookup("videos-v2-prod", "DummyV0Stub.write_video")
            list(f.starmap([(item, f"image_{item_id}.png", True) for item in items]))

            print("Videos took", time.time() - video_time)

            ok = True
        except Exception as e:
            import traceback

            finish_time = time.time()
            self.worker.notify_finish(
                item,
                {
                    "id": item_id,
                    "model": "chirp-v1-xl",
                    "n_audios": len(ids),
                    "ok": 0,
                    "gen_duration": finish_time - start_time,
                },
                queue_name="results:q",
            )

            traceback.print_exc()
            ok = False

            raise e

        finish_time = time.time()

        self.worker.notify_finish(
            item,
            {
                "id": item_id,
                "model": "chirp-v1-xl",
                "n_audios": len(ids),
                "ok": 1 if ok else 0,
                "gen_duration": finish_time - start_time,
            },
            queue_name="results:q",
        )

        return item.id


@stub.local_entrypoint()
def main():
    uid = str(uuid4())
    inputs = [
        json.dumps(
            dict(
                id=uid,
                # prompt_audio="5d9025f4-2158-4219-b9d6-abfbcc178db2.mp3",
                prompt_text="""
Yo, let me tell you 'bout my favorite dish
It's the sizzling, spicy fajitas that I wish
I chop up peppers and onions with glee
And grill that chicken 'til it's just right, you see

Taco taco in my belly-o
Taco taco in my belly-o
""",
                metadata={
                    "tags": "rap",
                    "is_square": True,
                },
            )
        )
    ]
    model = ChirpV1Stub()
    print(uid)
    for input in inputs:
        model.generate.remote(
            input,
        )