import os
import pathlib
import time
import json
from uuid import uuid4

import modal

from suno_utils.worker.generative_worker import ModelV1Worker

aws_secret = modal.Secret.from_name("studio-aws")


def download_model_wrapper():
    ModelV1Worker.download_models()
    from transformers import WavLMModel

    WavLMModel.from_pretrained("microsoft/wavlm-large")


def recursive_ls_dir(directory):
    paths = [os.path.join(root, file) for root, dirs, files in os.walk(directory) for file in files]
    paths += [os.path.join(root, dir) for root, dirs, files in os.walk(directory) for dir in dirs]
    print(paths)


image = (
    modal.Image.debian_slim()
    .apt_install("curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3")
    .run_commands(
        [
            'curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"',
            "unzip -q awscliv2.zip",
            "./aws/install",
        ]
    )
    .pip_install(
        "boto3",
        "transformers",
        "tokenizers",
        "encodec",
        "ctc_segmentation",
        "psutil",
        "redis",
        "gradio",
        "pydantic",
        "nnAudio",
    )
    .pip_install(
        "torch==2.1.0.dev20230531+cu118",
        "torchaudio==2.1.0.dev20230531+cu118",
        index_url="https://download.pytorch.org/whl/nightly/cu118",
    )
    .pip_install_from_pyproject(
        str(pathlib.Path(__file__).parent.parent.parent / "pyproject.toml"),
    )
    .run_function(download_model_wrapper, secret=aws_secret)
)
STUB_NAME = "bark-v1-ranking"
stub = modal.Stub(STUB_NAME, image=image)


@stub.cls(
    cpu=2.0,
    # memory=16384,
    gpu=modal.gpu.A10G(count=1),
    secret=aws_secret,
    timeout=300,
    container_idle_timeout=200,
    mounts=[
        modal.Mount.from_local_file(
            (pathlib.Path(__file__).parent / "assets/wave-bg.png").resolve(),
            remote_path="/suno/models/wave-bg.png",
        )
    ],
    keep_warm=1,
    concurrency_limit=4,  # don't want to spend more $ on bark
)
class ModelV1Stub:
    def __enter__(self):
        import torch

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

        from suno_utils.worker.generative_worker import ModelV1Worker

        self.worker = ModelV1Worker(0, bg_image="/suno/models/wave-bg.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))
        try:
            ok = self.worker.process_item(item)
        except Exception:
            import traceback

            traceback.print_exc()
            ok = False

        finish_time = time.time()

        self.worker.notify_finish(
            item,
            {
                "id": item.id,
                "ok": 1 if ok else 0,
                "gen_duration": finish_time - start_time,
            },
        )

        return item.id


@stub.cls(
    cpu=2.0,
    gpu=modal.gpu.A10G(count=1),
    secret=aws_secret,
    timeout=300,
    container_idle_timeout=200,
    mounts=[
        modal.Mount.from_local_file(
            (pathlib.Path(__file__).parent / "assets/wave-bg.png").resolve(),
            remote_path="/suno/models/wave-bg.png",
        )
    ],
    keep_warm=1,
    concurrency_limit=4,  # don't want to spend more $ on bark
)
class ModelV1UtilsStub:
    def __enter__(self):
        import torch

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

        from suno_utils.worker.generative_worker import ModelV1Worker

        self.worker = ModelV1Worker(0, bg_image="/suno/models/wave-bg.png")
        self.worker.preload_wavlm()

    @modal.method()
    def render_npz(self, id: str):
        self.worker.render_history_prompt(id)


@stub.local_entrypoint()
def main():
    [
        json.dumps(
            dict(
                id=str(uuid4()),
                # prompt_audio="5d9025f4-2158-4219-b9d6-abfbcc178db2.mp3",
                prompt_text="""My car is a clunker, it's really quite sad
But when I'm cruising down the street, I feel kinda rad
I might not have a Ferrari or a Lamborghini
But my ride gets me where I need to be, and that's all that matters to me""",
                metadata={},
            )
        )
    ]
    model = ModelV1UtilsStub()
    model.render_npz.remote("21f59001-6ff6-4aaa-be79-1a55a7238a79")


if __name__ == "__main__":
    import json
    from uuid import uuid4

    queue_item = dict(
        id=str(uuid4()),
        prompt_audio="5d9025f4-2158-4219-b9d6-abfbcc178db2.mp3",
        prompt_text="""My car is a clunker, it's really quite sad
But when I'm cruising down the street, I feel kinda rad
I might not have a Ferrari or a Lamborghini
But my ride gets me where I need to be, and that's all that matters to me""",
        metadata={},
    )

    print(queue_item)
    f = modal.Function.lookup(STUB_NAME, "ModelV1Stub.generate")
    f.spawn(json.dumps(queue_item))