# modal_wheel_builder.py
import os
from pathlib import Path

import modal
from suno_utils.worker.modal_base import get_modal_base_image_diffusion_with_flash_attention

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

# Define the container image with CUDA support and build dependencies
base_image = (
    modal.Image.from_registry("nvidia/cuda:12.8.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",
        ]
    )
    .pip_install("torch==2.7.0", "torchaudio==2.7.0", index_url="https://download.pytorch.org/whl/cu128")
    .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",
    )
    .add_local_python_source("suno_utils", copy=False)
)

base_image_with_flash_attention = (
    modal.Image.from_registry("nvidia/cuda:12.8.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",
        ]
    )
    .pip_install("torch==2.7.0", "torchaudio==2.7.0")
    .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("flash-attention")
)

# base_image = get_modal_base_image_diffusion_with_flash_attention()


APP_NAME = "wheel-builder"
# Create the Modal stub
app = modal.App(APP_NAME, image=base_image)


@app.function(
    image=base_image,
    cpu=64,
    memory=344064,
    volumes={"/wheels": model_store_volume},
    timeout=72000,
)
def build_flash_attention_wheel(cuda_version="12.8", torch_version="2.7.0", python_version="3.10"):
    """
    Build Flash Attention 3 as a wheel and save it to a Modal volume.

    Args:
        cuda_version (str): CUDA version (e.g., "12.4")
        torch_version (str): PyTorch version (e.g., "2.5.1")
        python_version (str): Python version (e.g., "3.10")

    Returns:
        str: Path to the built wheel in the volume
    """
    import subprocess
    import glob
    import shutil
    import sys

    # Debug Python version
    print(f"Python version in container: {sys.version}")
    print(f"Python executable: {sys.executable}")
    print(f"Python version info: {sys.version_info}")

    # Verify we're using the expected Python version
    actual_python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
    print(f"Expected Python version: {python_version}")
    print(f"Actual Python version: {actual_python_version}")

    if actual_python_version != python_version:
        print(
            f"WARNING: Python version mismatch! Expected {python_version}, got {actual_python_version}"
        )

    wheels_dir = Path("/wheels")

    # Create a unique directory for this build
    build_id = f"flash-attention-3-cuda{cuda_version}-torch{torch_version}-py{python_version}"
    build_dir = wheels_dir / build_id
    build_dir.mkdir(exist_ok=True)
    print("Installing build dependencies...")
    subprocess.run("pip install build wheel setuptools ninja", shell=True, check=True)

    # Clone the Flash Attention repository
    print("Cloning Flash Attention repository...")
    subprocess.run("git clone https://github.com/Dao-AILab/flash-attention.git", shell=True, check=True)

    os.chdir("flash-attention")

    # Checkout the specific commit before py_limited_api was introduced
    target_commit = "6ba57efea94c5a63cfd17d25a94e47b4065568a4"
    print(f"Checking out commit {target_commit} (before py_limited_api was added)...")
    subprocess.run(f"git checkout {target_commit}", shell=True, check=True)

    # Verify the commit
    result = subprocess.run("git rev-parse HEAD", shell=True, capture_output=True, text=True, check=True)
    current_commit = result.stdout.strip()
    print(f"Current commit: {current_commit}")

    if current_commit != target_commit:
        print(f"WARNING: Expected {target_commit}, got {current_commit}")
    else:
        print("✅ Successfully checked out the correct commit")

    os.chdir("hopper")

    # Check if there's a setup.py or pyproject.toml that might be setting the Python tags
    print("Checking build configuration files...")
    if os.path.exists("setup.py"):
        print("Found setup.py")
        with open("setup.py", "r") as f:
            setup_content = f.read()
            if "python_requires" in setup_content:
                print("setup.py contains python_requires")
            if "py_limited_api" in setup_content:
                print("⚠️ setup.py contains py_limited_api (this should NOT be present in this commit)")
            else:
                print("✅ No py_limited_api found - should build proper cp310 wheels")

    if os.path.exists("pyproject.toml"):
        print("Found pyproject.toml")
        with open("pyproject.toml", "r") as f:
            pyproject_content = f.read()
            print("pyproject.toml content preview:")
            print(pyproject_content[:500])

    # Try building with explicit Python tag control
    env = os.environ.copy()
    env["MAX_JOBS"] = "16"

    # Force the specific Python version tag
    env["SETUPTOOLS_SCM_PRETEND_VERSION"] = "3.0.0b1"

    print("Building Flash Attention wheel...")
    print("Note: Flash Attention 3 may use cp39-abi3 tags for cross-Python compatibility")

    # Use standard pip wheel build - let Flash Attention decide the tags
    subprocess.run(
        ["pip", "wheel", "-v", "--no-deps", "--wheel-dir", str(build_dir), "."], check=True, env=env
    )

    # Alternative modern build method (commented out, but available):
    # subprocess.run([
    #     "python", "-m", "build",
    #     "--wheel",
    #     "--outdir", str(build_dir)
    # ], check=True, env=env)

    # Find the built wheel
    wheel_files = glob.glob(str(build_dir / "*.whl"))
    if not wheel_files:
        # Fallback to dist directory
        wheel_files = glob.glob("dist/*.whl")

    if not wheel_files:
        raise FileNotFoundError("No wheel file found after build")

    wheel_file = wheel_files[0]
    wheel_filename = os.path.basename(wheel_file)

    print(f"Built wheel: {wheel_filename}")

    # Verify the wheel is compatible with current Python version
    print(f"Verifying wheel compatibility with Python {actual_python_version}...")
    try:
        # Install the wheel temporarily to test compatibility
        subprocess.run(["pip", "install", wheel_file], check=True)
        print("✅ Wheel is compatible with current Python version")

        # Try importing to verify it works
        try:
            import flash_attn_interface

            print("✅ Flash attention imported successfully")
        except ImportError as e:
            print(f"⚠️ Flash attention import failed: {e}")

        # Uninstall for cleanup
        subprocess.run(["pip", "uninstall", "-y", "flash-attn"], check=False)

    except subprocess.CalledProcessError as e:
        print(f"❌ Wheel compatibility test failed: {e}")

    # Copy wheel to mounted volume (if not already there)
    if not wheel_file.startswith(str(build_dir)):
        target_path = build_dir / wheel_filename
        shutil.copy(wheel_file, target_path)
        print(f"Copied wheel to: {target_path}")

    # List all wheels in the volume
    saved_wheels = os.listdir("/wheels")

    return {"hopper_wheel": wheel_filename, "all_saved_wheels": saved_wheels}


@app.function(
    image=base_image,
    cpu=16,
    memory=344064,
    gpu="H100",
    volumes={"/wheels": model_store_volume},
    timeout=72000,
)
def install_and_test_flash_attention_wheel(
    wheel_filename="flash_attn_3-3.0.0b1-cp310-cp310-linux_x86_64.whl",
):
    """
    Install the Flash Attention wheel from the volume and test it.

    Args:
        wheel_filename (str): The filename of the wheel in the volume

    Returns:
        dict: Test results and version info
    """
    import subprocess
    import os

    # Path to the wheel in the volume
    wheel_path = f"/wheels/flash-attention-3-cuda12.6-torch2.7.0-py3.10/flash_attn_3-3.0.0b1-cp310-cp310-linux_x86_64.whl"
    # wheel_path = f"/wheels/flash-attention-3-cuda12.4-torch2.5.1-py3.10/{wheel_filename}"

    # Check if wheel exists
    if not os.path.exists(wheel_path):
        available_wheels = os.listdir("/wheels")
        return {"error": f"Wheel {wheel_filename} not found", "available_wheels": available_wheels}

    print(f"Installing wheel from: {wheel_path}")

    # Install the wheel
    try:
        subprocess.run(["pip", "install", wheel_path], check=True, capture_output=True, text=True)
        print("Wheel installed successfully!")
    except subprocess.CalledProcessError as e:
        return {"error": "Failed to install wheel", "stdout": e.stdout, "stderr": e.stderr}

    # Test the installation
    try:
        print("Importing flash_attn")
        import time
        # while True:
        #     time.sleep(1)

        # # Try to get version info
        # version = getattr(flash_attn_interface, "__version__", "unknown")
        version = "3.0.0b1"  # Set default version for the wheel

        # Try a basic functionality test
        import torch

        if torch.cuda.is_available():
            print(f"CUDA available: {torch.cuda.is_available()}")
            print(f"CUDA device count: {torch.cuda.device_count()}")

            # Test basic flash attention functionality
            try:
                from flash_attn_interface import flash_attn_with_kvcache

                print("flash_attn_with_kvcache imported successfully")

                # Create small test tensors
                batch_size, seq_len, num_heads, head_dim = 1, 16, 8, 64
                q = torch.randn(
                    batch_size, seq_len, num_heads, head_dim, dtype=torch.float16, device="cuda"
                )
                k = torch.randn(
                    batch_size, seq_len, num_heads, head_dim, dtype=torch.float16, device="cuda"
                )
                v = torch.randn(
                    batch_size, seq_len, num_heads, head_dim, dtype=torch.float16, device="cuda"
                )

                # Run flash attention
                output = flash_attn_with_kvcache(q, k, v)
                print(f"Flash attention test successful! Output shape: {output.shape}")

                return {
                    "success": True,
                    "version": version,
                    "cuda_available": True,
                    "cuda_device_count": torch.cuda.device_count(),
                    "test_output_shape": list(output.shape),
                }

            except Exception as e:
                return {
                    "success": True,
                    "version": version,
                    "cuda_available": True,
                    "cuda_device_count": torch.cuda.device_count(),
                    "flash_attn_test_error": str(e),
                }
        else:
            return {
                "success": True,
                "version": version,
                "cuda_available": False,
                "note": "CUDA not available for functional testing",
            }

    except ImportError as e:
        return {"error": "Failed to import flash_attn after installation", "import_error": str(e)}
    except Exception as e:
        return {"error": "Failed to install and test flash_attn", "exception": str(e)}


@app.function(
    image=base_image_with_flash_attention,  # Use the image that already has flash-attention
    cpu=4,
    memory=8192,
    timeout=300,
)
def test_flash_attention_wheel():
    """Test the flash-attention installation that's already in the image"""
    try:
        import flash_attn

        print(f"flash_attn: {flash_attn}")

        version = getattr(flash_attn, "__version__", "unknown")
        print(f"Flash Attention version: {version}")

        import torch

        print(f"CUDA available: {torch.cuda.is_available()}")

        if torch.cuda.is_available():
            from flash_attn import flash_attn_func

            print("flash_attn_func imported successfully")

        return {"success": True, "version": version, "cuda_available": torch.cuda.is_available()}

    except ImportError as e:
        return {"error": "Failed to import flash_attn", "import_error": str(e)}


@app.local_entrypoint()
def main():
    # Test installing and using the wheel from the volume
    print("Testing wheel installation from volume...")
    test_result = install_and_test_flash_attention_wheel.remote()
    print(f"Test result: {test_result}")

    # Also test the pre-installed version
    print("\nTesting pre-installed flash-attention...")
    preinstalled_result = test_flash_attention_wheel.remote()
    print(f"Pre-installed test result: {preinstalled_result}")