import os
import tempfile
import time

import ffmpeg
import modal
import numpy as np

from suno_utils.tasks.gpt.generation import *
from suno_utils.tasks.gpt_v2 import chirp_v1
from suno_utils.worker.loader import S3Loader
from suno_utils.worker.modal_base import MODAL_MOUNTS, get_modal_base_image
from suno_utils.worker.schema import QueueItem
from suno_utils.worker.utils import download_models_to_dir

CHIRP_V1_MODELS = dict(
    codec_path="georg/trained_models/chirp_v1/codec.pt",
)

MOUNT_PATH = "/suno/models"


class ConcatWorker(S3Loader):
    def __init__(self, bg_image=None):
        super().__init__(bg_image=bg_image)

        self.bg_image = bg_image

    def preload(self):
        start_time = time.time()

        chirp_v1.preload_codec_models(
            f"{MOUNT_PATH}/{CHIRP_V1_MODELS['codec_path']}",
        )

        finish_time = time.time()
        print(f"Preloading took {finish_time - start_time}s")

    @staticmethod
    def download_models(dir_path=MOUNT_PATH):
        """Use AWS CLI to download models if they don't exist."""
        files = list(CHIRP_V1_MODELS.values())
        download_models_to_dir(files, dir_path)

    def concat_videos(self, new_id: str, ids: list[str]):
        from suno_utils.worker.settings import s3_client

        print(ids)

        filenames = [f"{id}.mp4" for id in ids]

        with tempfile.TemporaryDirectory() as td:
            full_paths = []
            for fname in filenames:
                full_path = os.path.join(td, fname)
                full_paths.append(full_path)
                with open(full_path, "wb") as tmp_file:
                    s3_client.download_fileobj(
                        "suno-data-uploads",
                        f"studio/uploads/{fname}",
                        tmp_file,
                    )

                with open(full_path.replace(".mp4", ".npz"), "wb") as tmp_file:
                    s3_client.download_fileobj(
                        "suno-data-uploads",
                        f"studio/uploads/{fname.replace('.mp4', '.npz')}",
                        tmp_file,
                    )

            npzs = [np.load(p.replace(".mp4", ".npz")) for p in full_paths]

            audio = chirp_v1.codec_decode(
                np.concatenate([n["v1_raw"] if "v1_raw" in n else n["v1.0_raw"] for n in npzs], axis=0)[
                    ::3, -8:
                ]
            )
            # we keep wav operation here for high resolution, but we don't actually save these wavs
            full_audio_path = os.path.join(td, "concat.wav")
            audio.write_wav(full_audio_path)

            inputs = [(ffmpeg.input(p), ffmpeg.input(p.replace(".mp4", ".wav"))) for p in full_paths]

            full_path = os.path.join(td, "concat.mp4")

            joined = ffmpeg.concat(*[x for i in inputs for x in (i[0]["v"],)], v=1, a=0)
            ffmpeg.output(joined, full_path).run()
            full_path_2 = os.path.join(td, "concat_2.mp4")

            ffmpeg.concat(ffmpeg.input(full_path), ffmpeg.input(full_audio_path), v=1, a=1).output(
                full_path_2
            ).run()

            s3_client.upload_file(full_path_2, "suno-data-uploads", f"studio/uploads/{new_id}.mp4")


def download_model_wrapper_5():
    ConcatWorker.download_models()


image = get_modal_base_image().run_function(
    download_model_wrapper_5, secret=modal.Secret.from_name("studio-aws")
)
STUB_NAME = "concat-v0-alpha"
stub = modal.Stub(STUB_NAME, image=image)


@stub.cls(
    cpu=2.0,
    gpu="any",
    secrets=[
        modal.Secret.from_name("studio-aws"),
        modal.Secret.from_dict({"SUNO_ASSETS_PATH": "/suno/models/assets"}),
    ],
    timeout=60,
    container_idle_timeout=60,
    mounts=MODAL_MOUNTS,
    concurrency_limit=8,
)
class ConcatStub:
    def __enter__(self):
        self.worker = ConcatWorker(bg_image="/suno/models/assets/wave-bg-3.png")
        self.worker.preload()

    @modal.method()
    def concat_videos(self, new_id: str, ids: list[str], callback_url: str | None = None):
        self.worker.concat_videos(
            new_id,
            ids,
        )

        if callback_url:
            self.worker.notify_finish(
                QueueItem(id=new_id, metadata={}, callback_url=callback_url),
                {
                    "id": new_id,
                    "type": "concat_videos",
                },
            )


@stub.local_entrypoint()
def main():
    import uuid

    new_id = str(uuid.uuid4())
    print(new_id)
    model = ConcatStub()
    model.concat_videos.remote(
        new_id,
        [
            "7d3dba59-2ab5-4d2d-9d05-69b298820edf_1",
            "162b5c15-4e56-4d49-8402-1a48c5b231af_1",
        ],
    )