import pathlib

import modal

MOUNT_PATH = "/suno/models"

# Define the Modal volume to access the wheels
wheels_volume = modal.Volume.from_name("suno-wheels", create_if_missing=False)


def get_modal_base_image():
    return (
        modal.Image.debian_slim(python_version="3.10")
        .apt_install(
            "curl",
            "unzip",
        )
        .run_commands(
            [
                'curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"',
                "unzip -q awscliv2.zip",
                "./aws/install",
            ]
        )
        .apt_install(
            "ffmpeg",
            "libsox-fmt-mp3",
            "sox",
        )
        .pip_install("torch==2.5.1")  # this is cause flash-attn can't work with 2.5 yet
        .pip_install_private_repos(
            "github.com/suno-ai/glockenspiel.git@d9d6e594dc#subdirectory=descript-audio-codec&egg=descript-audio-codec",
            git_user="mcamac",
            secrets=[modal.Secret.from_name("victor-modal-github-token")],
        )
        .pip_install(
            "boto3",
            "tokenizers",
            "encodec",
            "ctc_segmentation",
            "psutil",
            "redis",
            "pydantic",
            "nnAudio",
        )
        .pip_install_from_pyproject(
            str(pathlib.Path(__file__).parent.parent.parent / "pyproject.toml"),
        )
        .dockerfile_commands(
            [
                "COPY --from=datadog/serverless-init /datadog-init /app/datadog-init",
                'ENTRYPOINT ["/app/datadog-init"]',
            ]
        )
        .apt_install(
            "libogg0",
            "libopus0",
            "opus-tools",
        )
    )


def install_flash_attention_from_wheel_gpt():
    """Install flash-attention from pre-built wheel in the volume"""
    import subprocess
    import os

    wheel_path = "/wheels/flash-attention-3-cuda12.4-torch2.5.1-py3.10/flash_attn_3-3.0.0b1-cp310-cp310-linux_x86_64.whl"

    # Check if wheel exists
    if not os.path.exists(wheel_path):
        available_wheels = os.listdir("/wheels")
        print(f"Wheel not found at {wheel_path}")
        print(f"Available files: {available_wheels}")
        raise FileNotFoundError(f"Wheel not found at {wheel_path}")

    print(f"Installing flash-attention from wheel: {wheel_path}")
    subprocess.run(["pip", "install", wheel_path], check=True)
    print("flash-attention installed successfully from wheel")


def install_flash_attention_from_wheel_diffusion():
    """Install flash-attention from pre-built wheel in the volume"""
    import subprocess
    import os

    wheel_path = "/wheels/flash-attention-3-cuda12.6-torch2.7.0-py3.10/flash_attn_3-3.0.0b1-cp310-cp310-linux_x86_64.whl"

    # Check if wheel exists
    if not os.path.exists(wheel_path):
        available_wheels = os.listdir("/wheels")
        print(f"Wheel not found at {wheel_path}")
        print(f"Available files: {available_wheels}")
        raise FileNotFoundError(f"Wheel not found at {wheel_path}")

    print(f"Installing flash-attention from wheel: {wheel_path}")
    subprocess.run(["pip", "install", wheel_path], check=True)
    print("flash-attention installed successfully from wheel")


def install_flash_attention():
    import subprocess

    subprocess.run(
        "git clone https://github.com/Dao-AILab/flash-attention.git",
        shell=True,
    )
    subprocess.run(
        "cd flash-attention/hopper && MAX_JOBS=8 python setup.py install",
        shell=True,
    )
    print("flash-attention installed successfully")


def get_modal_base_image_diffusion_with_flash_attention():
    # Pytorch 2.7.0 + CUDA 12.8
    # This image is used for diffusion models, it has no flash infer
    base_image = (
        modal.Image.from_registry("nvidia/cuda:12.8.0-devel-ubuntu24.04", add_python="3.10")
        .apt_install("curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3", "zlib1g-dev", "git", "clang")
        .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.7.0", "torchaudio==2.7.0")
        # .pip_install("flashinfer-python", index_url="https://flashinfer.ai/whl/cu124/torch2.5/")
        .pip_install_private_repos(
            "github.com/suno-ai/glockenspiel.git@a5dba4e50#subdirectory=descript-audio-codec&egg=descript-audio-codec",
            git_user="mcamac",
            secrets=[modal.Secret.from_name("victor-modal-github-token")],
        )
        .pip_install_private_repos(
            "github.com/suno-ai/neon.git@1c83548#subdirectory=hoot",
            git_user="mcamac",
            secrets=[modal.Secret.from_name("victor-modal-github-token")],
        )
        .pip_install(
            "boto3",
            "transformers",
            "tokenizers",
            "encodec",
            "ctc_segmentation",
            "psutil",
            "redis",
            "pydantic",
            "nnAudio",
            "rpyc",
            "biopython>=1.81",  # TODO: don't love this depdendency, for hoot
            "pynvml",  # for torch cuda utilization
            "torchsde",
            "ninja",
            "wheel",
        )
        .pip_install_from_pyproject(
            str(pathlib.Path(__file__).parent.parent.parent / "pyproject.toml"),
        )
        .pip_install("torch==2.7.0", "torchaudio==2.7.0")
        .run_function(
            install_flash_attention_from_wheel_diffusion,
            cpu=16,
            memory=1024 * 20,
            timeout=1800,
            volumes={"/wheels": wheels_volume},
        )
        .apt_install(
            "libogg0",
            "libopus0",
            "opus-tools",
        )
    )
    return base_image


def get_modal_base_image_with_flash_attention():
    # Be very careful with this image!
    # for example, flash infer version is NOT fixed.
    # We need to check performance, not only infer speed, but also compile / cold start time.
    base_image = (
        modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.10")
        .apt_install("curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3", "zlib1g-dev", "git", "clang")
        .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.5.1", "torchaudio==2.5.1")
        .pip_install("flashinfer-python", index_url="https://flashinfer.ai/whl/cu124/torch2.5/")
        .pip_install_private_repos(
            "github.com/suno-ai/glockenspiel.git@a5dba4e50#subdirectory=descript-audio-codec&egg=descript-audio-codec",
            git_user="mcamac",
            secrets=[modal.Secret.from_name("victor-modal-github-token")],
        )
        .pip_install_private_repos(
            "github.com/suno-ai/neon.git@1c83548#subdirectory=hoot",
            git_user="mcamac",
            secrets=[modal.Secret.from_name("victor-modal-github-token")],
        )
        .pip_install(
            "boto3",
            "transformers",
            "tokenizers",
            "encodec",
            "ctc_segmentation",
            "psutil",
            "redis",
            "pydantic",
            "nnAudio",
            "rpyc",
            "biopython>=1.81",  # TODO: don't love this depdendency, for hoot
            "pynvml",  # for torch cuda utilization
            "torchsde",
            "ninja",
            "wheel",
        )
        .pip_install_from_pyproject(
            str(pathlib.Path(__file__).parent.parent.parent / "pyproject.toml"),
        )
        .apt_install(
            "libogg0",
            "libopus0",
            "opus-tools",
        )
        .pip_install_from_pyproject(
            str(pathlib.Path(__file__).parent.parent.parent / "pyproject.toml"),
        )
        .run_function(
            install_flash_attention_from_wheel_gpt,
            cpu=16,
            memory=1024 * 20,
            timeout=1800,
            volumes={"/wheels": wheels_volume},
        )
    )
    return base_image