#!/usr/bin/env python3
"""
Setup Suno Environment with Flash Attention v3 (Hopper) and v2 (fixed commit)
This version installs FA3 first (for H100 optimization) then FA2 for compatibility
FA3 is compatible with FA2, so both can coexist
"""

import os
import sys
import subprocess
import shutil
from pathlib import Path
from datetime import datetime
from typing import Optional, Tuple


class Colors:
    """ANSI color codes for terminal output"""
    RED = '\033[0;31m'
    GREEN = '\033[0;32m'
    YELLOW = '\033[1;33m'
    BLUE = '\033[0;34m'
    NC = '\033[0m'  # No Color


class SunoEnvSetup:
    """Setup Suno environment with Flash Attention v3 + v2"""

    def __init__(self, env_name: str = "suno_env_fa2_fa3"):
        self.env_name = env_name
        self.script_dir = Path(__file__).parent.resolve()
        self.tmp_dir = self.script_dir / "tmp"
        self.tmp_dir.mkdir(exist_ok=True)

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_file = self.tmp_dir / f"setup_{env_name}_{timestamp}.log"

        # Paths to be set during setup
        self.suno_utils_path: Optional[Path] = None
        self.flash_attn_v3_path: Optional[Path] = None
        self.flash_attn_path: Optional[Path] = None
        self.flash_attn_commit = "c5b0c631074e4c8d53fdebea2d71ea621baf9344"

    def log(self, message: str):
        """Log message to console and file"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        formatted = f"{Colors.GREEN}[{timestamp}]{Colors.NC} {message}"
        print(formatted)
        with open(self.log_file, 'a') as f:
            # Remove color codes for log file
            clean_msg = message.replace(Colors.RED, '').replace(Colors.GREEN, '') \
                              .replace(Colors.YELLOW, '').replace(Colors.BLUE, '') \
                              .replace(Colors.NC, '')
            f.write(f"[{timestamp}] {clean_msg}\n")

    def error(self, message: str):
        """Log error and exit"""
        print(f"{Colors.RED}[ERROR]{Colors.NC} {message}")
        with open(self.log_file, 'a') as f:
            f.write(f"[ERROR] {message}\n")
        sys.exit(1)

    def warning(self, message: str):
        """Log warning"""
        print(f"{Colors.YELLOW}[WARNING]{Colors.NC} {message}")
        with open(self.log_file, 'a') as f:
            f.write(f"[WARNING] {message}\n")

    def info(self, message: str):
        """Log info"""
        print(f"{Colors.BLUE}[INFO]{Colors.NC} {message}")
        with open(self.log_file, 'a') as f:
            f.write(f"[INFO] {message}\n")

    def run_command(self, cmd: list[str], check: bool = True,
                   capture_output: bool = True) -> subprocess.CompletedProcess:
        """Run shell command and log output"""
        cmd_str = ' '.join(cmd)
        self.log(f"Running: {cmd_str}")

        try:
            result = subprocess.run(
                cmd,
                check=check,
                capture_output=capture_output,
                text=True
            )

            if capture_output and result.stdout:
                with open(self.log_file, 'a') as f:
                    f.write(result.stdout)
            if capture_output and result.stderr:
                with open(self.log_file, 'a') as f:
                    f.write(result.stderr)

            return result
        except subprocess.CalledProcessError as e:
            self.error(f"Command failed: {cmd_str}\n{e.stderr if e.stderr else ''}")

    def setup_suno_utils_path(self):
        """Setup Suno Utils path"""
        default_path = Path.home() / "projects" / "glockenspiel" / "suno_utils"

        print(f"{Colors.BLUE}Suno Utils Setup{Colors.NC}")
        print("Enter the path to suno_utils package")
        print(f"Press Enter to use default: {default_path}")

        user_input = input("Path: ").strip()

        if not user_input:
            self.suno_utils_path = default_path
            self.log(f"Using default suno_utils path: {self.suno_utils_path}")
        else:
            self.suno_utils_path = Path(user_input).expanduser()
            self.log(f"Using custom suno_utils path: {self.suno_utils_path}")

        if not self.suno_utils_path.exists():
            self.error(f"suno_utils not found at {self.suno_utils_path}")

        self.log(f"suno_utils found at {self.suno_utils_path}")

        # Check if it's a valid Python package
        if not (self.suno_utils_path / "setup.py").exists() and \
           not (self.suno_utils_path / "pyproject.toml").exists():
            self.warning(f"{self.suno_utils_path} doesn't appear to be a valid Python package")
            self.warning("Missing setup.py or pyproject.toml")
            response = input("Continue anyway? (y/n): ").strip().lower()
            if response != 'y':
                self.error("Valid suno_utils package required")

    def setup_flash_attention_v3_path(self):
        """Setup Flash Attention v3 path (main repository)"""
        default_path = Path.home() / "projects" / "flash-attention"

        print(f"{Colors.BLUE}Flash Attention v3 Setup{Colors.NC}")
        print("Enter the path to Flash Attention repository")
        print(f"Press Enter to use default: {default_path}")

        user_input = input("Path: ").strip()

        if not user_input:
            self.flash_attn_v3_path = default_path
            self.log(f"Using default Flash Attention v3 path: {self.flash_attn_v3_path}")
        else:
            self.flash_attn_v3_path = Path(user_input).expanduser()
            self.log(f"Using custom Flash Attention v3 path: {self.flash_attn_v3_path}")

        if not self.flash_attn_v3_path.exists():
            self.warning(f"Flash Attention not found at {self.flash_attn_v3_path}")
            response = input("Would you like to clone Flash Attention to this location? (y/n): ").strip().lower()

            if response == 'y':
                self.log("Cloning Flash Attention repository...")
                self.flash_attn_v3_path.parent.mkdir(parents=True, exist_ok=True)

                self.run_command([
                    "git", "clone",
                    "https://github.com/Dao-AILab/flash-attention.git",
                    str(self.flash_attn_v3_path)
                ])

                # Initialize submodules
                self.log("Initializing Flash Attention submodules...")
                self.run_command(
                    ["git", "submodule", "update", "--init", "--recursive"],
                    check=False
                )

                self.log("Flash Attention cloned successfully!")
            else:
                self.error("Flash Attention repository required")
        else:
            self.log(f"Flash Attention repository found at {self.flash_attn_v3_path}")
            response = input("Would you like to update the Flash Attention repository? (y/n): ").strip().lower()

            if response == 'y':
                self.log("Updating Flash Attention repository...")
                original_dir = os.getcwd()
                try:
                    os.chdir(self.flash_attn_v3_path)
                    self.run_command(["git", "pull"], check=False)
                    self.run_command(["git", "submodule", "update", "--init", "--recursive"], check=False)
                finally:
                    os.chdir(original_dir)

        # Check if hopper directory exists
        hopper_path = self.flash_attn_v3_path / "hopper"
        if not hopper_path.exists():
            self.warning(f"Flash Attention v3 hopper directory not found at {hopper_path}")
            self.warning("This might be an older version of the repository")
            self.warning("Flash Attention v3 requires the hopper subdirectory for H100 optimization")

    def setup_flash_attention_v2_path(self):
        """Setup Flash Attention v2 path (fixed commit)"""
        self.flash_attn_path = Path(f"/tmp/flash-attention-{self.env_name}")

        self.log(f"Cloning Flash Attention v2 to {self.flash_attn_path} (commit {self.flash_attn_commit})...")

        # Remove if exists
        if self.flash_attn_path.exists():
            shutil.rmtree(self.flash_attn_path)

        # Clone repository
        self.run_command([
            "git", "clone",
            "https://github.com/Dao-AILab/flash-attention.git",
            str(self.flash_attn_path)
        ])

        # Checkout specific commit
        original_dir = os.getcwd()
        try:
            os.chdir(self.flash_attn_path)
            self.log(f"Checking out commit {self.flash_attn_commit}...")
            self.run_command(["git", "checkout", self.flash_attn_commit])

            # Initialize submodules
            self.log("Initializing Flash Attention v2 submodules...")
            self.run_command(["git", "submodule", "update", "--init", "--recursive"], check=False)
        finally:
            os.chdir(original_dir)

        self.log(f"Flash Attention v2 cloned successfully at commit {self.flash_attn_commit}!")

    def check_cuda_version(self):
        """Check CUDA version"""
        self.log("Checking CUDA version...")
        try:
            result = self.run_command(["nvcc", "--version"], check=False)
            if result.returncode == 0:
                # Parse CUDA version from output
                for line in result.stdout.split('\n'):
                    if 'release' in line.lower():
                        self.log(f"CUDA version: {line}")
                        # Check if CUDA >= 12.3
                        # This is a simplified check
                        if 'release 12.' in line.lower() or 'release 13.' in line.lower():
                            parts = line.split('release ')[1].split(',')[0].split('.')
                            major, minor = int(parts[0]), int(parts[1]) if len(parts) > 1 else 0
                            if major < 12 or (major == 12 and minor < 3):
                                self.warning(f"Flash Attention v3 requires CUDA >= 12.3")
                                self.warning("Installation may fail or performance may be suboptimal")
                        break
        except FileNotFoundError:
            self.warning("nvcc not found. Unable to check CUDA version.")
            self.warning("Flash Attention v3 requires CUDA >= 12.3")

    def check_conda_env_exists(self) -> bool:
        """Check if conda environment exists"""
        result = self.run_command(["conda", "env", "list"], check=False)
        for line in result.stdout.split('\n'):
            if line.strip().startswith(self.env_name + ' '):
                return True
        return False

    def get_conda_base(self) -> Path:
        """Get conda base directory"""
        result = self.run_command(["conda", "info", "--base"])
        return Path(result.stdout.strip())

    def run_in_conda_env(self, cmd: list[str], check: bool = True) -> subprocess.CompletedProcess:
        """Run command in conda environment"""
        conda_base = self.get_conda_base()
        conda_sh = conda_base / "etc" / "profile.d" / "conda.sh"

        full_cmd = f"source {conda_sh} && conda activate {self.env_name} && {' '.join(cmd)}"

        return subprocess.run(
            ["bash", "-c", full_cmd],
            check=check,
            capture_output=True,
            text=True
        )

    def setup(self):
        """Main setup function"""
        self.log("Starting Suno Environment Setup with Flash Attention v3 + v2")
        self.log(f"Environment name: {self.env_name}")
        self.log(f"Log file: {self.log_file}")
        self.log("Strategy: Install FA3 (H100 optimized) first, then FA2 (fixed commit) for compatibility")

        # Check CUDA version
        self.check_cuda_version()

        # Setup paths
        self.setup_suno_utils_path()
        self.setup_flash_attention_v3_path()
        self.setup_flash_attention_v2_path()

        # Check if environment exists
        if self.check_conda_env_exists():
            self.warning(f"Environment {self.env_name} already exists.")
            response = input("Do you want to remove and recreate it? (y/n): ").strip().lower()
            if response == 'y':
                self.log("Removing existing environment...")
                self.run_command(["conda", "env", "remove", "-n", self.env_name, "-y"])
            else:
                self.error("Environment already exists. Exiting.")

        # Step 1: Create conda environment
        self.log("Step 1: Creating conda environment with Python 3.10...")
        self.run_command(["conda", "create", "-n", self.env_name, "python=3.10.15", "-y"])
        self.log("Conda environment created!")

        # Step 2: Install suno_utils
        self.log(f"Step 2: Installing suno_utils from {self.suno_utils_path}...")
        self.run_in_conda_env(["pip", "install", "-e", str(self.suno_utils_path)])
        self.log("suno_utils installed!")

        # Step 3: Check torch version
        self.log("Step 3: Checking torch version installed by suno_utils...")
        result = self.run_in_conda_env(
            ["python", "-c", "import torch; print(torch.__version__)"],
            check=False
        )
        torch_version = result.stdout.strip() if result.returncode == 0 else "None"
        self.log(f"Torch version installed by suno_utils: {torch_version}")

        # Step 4: Uninstall torch
        self.log("Step 4: Uninstalling torch installed by suno_utils...")
        self.run_in_conda_env(["pip", "uninstall", "torch", "torchvision", "torchaudio", "-y"])
        self.log("Torch uninstalled!")

        # Step 5: Install PyTorch 2.6.0
        self.log("Step 5: Installing PyTorch 2.6.0 with CUDA 12.4 support...")
        self.run_in_conda_env([
            "pip", "install",
            "torch==2.6.0", "torchvision==0.21.0", "torchaudio==2.6.0",
            "--index-url", "https://download.pytorch.org/whl/cu124"
        ])
        self.log("PyTorch installed with CUDA support!")

        # Verify PyTorch
        result = self.run_in_conda_env(["python", "-c", "import torch; print(f'PyTorch version: {torch.__version__}')"])
        print(result.stdout)
        result = self.run_in_conda_env(["python", "-c", "import torch; print(f'CUDA available: {torch.cuda.is_available()}')"])
        print(result.stdout)

        # Step 6: Build Flash Attention v3
        hopper_path = self.flash_attn_v3_path / "hopper"

        if not hopper_path.exists():
            self.warning(f"Flash Attention v3 hopper directory not found at {hopper_path}")
            self.warning("Skipping FA3 installation, will only install FA2")
        else:
            self.log(f"Step 6: Building Flash Attention v3 (H100 optimized) from {hopper_path}...")

            original_dir = os.getcwd()
            try:
                os.chdir(hopper_path)

                # Clean builds
                self.log("Cleaning previous Flash Attention v3 builds...")
                for path in ["build", "dist"]:
                    if Path(path).exists():
                        shutil.rmtree(path)
                for path in hopper_path.glob("*.egg-info"):
                    shutil.rmtree(path)

                # Install dependencies
                self.log("Installing build dependencies...")
                self.run_in_conda_env(["pip", "install", "packaging", "ninja"])

                # Build FA3
                self.log("Building and installing Flash Attention v3 (this may take 10-30 minutes)...")
                self.log("Note: Flash Attention v3 is optimized for H100/H800 GPUs")

                env = os.environ.copy()
                env["MAX_JOBS"] = "4"
                env["FLASH_ATTENTION_FORCE_BUILD"] = "TRUE"

                conda_base = self.get_conda_base()
                conda_sh = conda_base / "etc" / "profile.d" / "conda.sh"

                cmd = f"source {conda_sh} && conda activate {self.env_name} && python setup.py install"
                result = subprocess.run(
                    ["bash", "-c", cmd],
                    env=env,
                    capture_output=True,
                    text=True
                )

                if result.returncode != 0:
                    self.warning("Flash Attention v3 installation failed")
                    self.warning("This is expected if not running on H100/H800")
                    self.warning("Continuing with FA2 installation...")

                # Test FA3
                self.log("Testing Flash Attention v3 installation...")
                test_code = """
try:
    import flash_attn_interface
    print('✓ Flash Attention v3 successfully installed')
    print(f'  Module location: {flash_attn_interface.__file__}')
except ImportError as e:
    print('✗ Flash Attention v3 not available')
    print(f'  Error: {e}')
"""
                result = self.run_in_conda_env(["python", "-c", test_code], check=False)
                print(result.stdout)

            finally:
                os.chdir(original_dir)

        # Step 7: Build Flash Attention v2
        self.log(f"Step 7: Building Flash Attention v2 from {self.flash_attn_path} (commit {self.flash_attn_commit})...")

        original_dir = os.getcwd()
        try:
            os.chdir(self.flash_attn_path)

            # Clean builds
            self.log("Cleaning previous Flash Attention v2 builds...")
            for path in ["build", "dist"]:
                if Path(path).exists():
                    shutil.rmtree(path)
            for path in self.flash_attn_path.glob("*.egg-info"):
                shutil.rmtree(path)

            # Install ninja
            self.log("Ensuring ninja is installed...")
            self.run_in_conda_env(["pip", "install", "ninja"])

            # Build FA2
            self.log("Building and installing Flash Attention v2 (this may take 10-30 minutes)...")

            env = os.environ.copy()
            env["MAX_JOBS"] = "4"
            env["FLASH_ATTENTION_FORCE_BUILD"] = "TRUE"

            conda_base = self.get_conda_base()
            conda_sh = conda_base / "etc" / "profile.d" / "conda.sh"

            cmd = f"source {conda_sh} && conda activate {self.env_name} && pip install . --no-build-isolation"
            result = subprocess.run(
                ["bash", "-c", cmd],
                env=env,
                capture_output=True,
                text=True,
                cwd=str(self.flash_attn_path)
            )

            if result.returncode != 0:
                self.error("Flash Attention v2 installation failed")

            # Test FA2
            self.log("Testing Flash Attention v2 installation...")
            test_code = """
try:
    import flash_attn
    print('✓ Flash Attention v2 successfully installed')
    print(f'  Version: {flash_attn.__version__}')
    print(f'  Module location: {flash_attn.__file__}')
except ImportError as e:
    print('✗ Flash Attention v2 not available')
    print(f'  Error: {e}')
"""
            result = self.run_in_conda_env(["python", "-c", test_code], check=False)
            print(result.stdout)

        finally:
            os.chdir(original_dir)

        # Step 8: Install additional packages
        self.log("Step 8: Installing additional Python packages...")
        packages = [
            "wandb", "nnAudio", "deepspeed", "auraloss", "torchsde", "g2p_en",
            "transformers==4.57.0", "modal==1.0.1", "better_profanity", "encodec",
            "pytorch-lightning", "sentencepiece", "tiktoken", "einops", "ffmpeg-python"
        ]
        self.run_in_conda_env(["pip", "install"] + packages)
        self.log("Additional packages installed!")

        # Step 9: Install sox
        self.log("Step 9: Installing sox via conda...")
        self.run_command(["conda", "install", "-n", self.env_name, "-c", "conda-forge", "sox", "-y"])
        self.log("Sox installed!")

        # Verification
        self.log("Running verification tests...")
        verify_script = self.script_dir / "verify_imports.py"
        if verify_script.exists():
            result = self.run_in_conda_env(["python", str(verify_script)], check=False)
            print(result.stdout)
        else:
            self.warning("verify_imports.py not found, skipping verification")

        # Create activation script
        self.log("Creating activation script...")
        activation_script = self.script_dir / f"activate_{self.env_name}.sh"

        conda_base = self.get_conda_base()

        with open(activation_script, 'w') as f:
            f.write(f"""#!/bin/bash
# Activation script for {self.env_name} environment with Flash Attention v3 + v2

# Source conda
source "{conda_base}/etc/profile.d/conda.sh"

# Activate environment
conda activate {self.env_name}

# Set environment variables
export CUDA_VISIBLE_DEVICES=${{CUDA_VISIBLE_DEVICES:-0}}
export OMP_NUM_THREADS=1
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_$USER
export PYTHONPATH="{self.flash_attn_v3_path}/hopper:$PYTHONPATH"

echo "Environment {self.env_name} activated (Flash Attention v3 + v2)!"
echo "Python: $(which python)"
echo "PyTorch: $(python -c 'import torch; print(torch.__version__)' 2>/dev/null || echo 'Not available')"
echo "CUDA available: $(python -c 'import torch; print(torch.cuda.is_available())' 2>/dev/null || echo 'Unknown')"
echo "Flash Attention v3 path: {self.flash_attn_v3_path}"
echo "Flash Attention v2 path: {self.flash_attn_path}"
echo "Suno Utils path: {self.suno_utils_path}"

# Test Flash Attention versions
echo ""
echo "Flash Attention Status:"
python -c "
try:
    import flash_attn_interface
    print('  ✓ FA3 (H100 optimized): Available')
except ImportError:
    print('  ✗ FA3: Not available')

try:
    import flash_attn
    print('  ✓ FA2 (commit {self.flash_attn_commit}): Available')
except ImportError:
    print('  ✗ FA2: Not available')
" 2>/dev/null
""")

        activation_script.chmod(0o755)

        # Final summary
        self.log("=" * 48)
        self.log("Environment setup completed!")
        self.log(f"Flash Attention v3 path: {self.flash_attn_v3_path}")
        self.log(f"Flash Attention v2 path: {self.flash_attn_path}")
        self.log(f"Flash Attention v2 commit: {self.flash_attn_commit}")
        self.log(f"Suno Utils path: {self.suno_utils_path}")
        self.log("")
        self.log("To activate the environment, run:")
        self.log(f"  source {activation_script}")
        self.log("Or:")
        self.log(f"  conda activate {self.env_name}")
        self.log("")
        self.log("Note: This environment has both FA3 (H100 optimized) and FA2 (compatibility)")
        self.log("=" * 48)


def main():
    """Main entry point"""
    env_name = sys.argv[1] if len(sys.argv) > 1 else "suno_env_fa2_fa3"

    setup = SunoEnvSetup(env_name)

    try:
        setup.setup()
    except KeyboardInterrupt:
        setup.error("Setup interrupted by user")
    except Exception as e:
        setup.error(f"Setup failed with error: {e}")


if __name__ == "__main__":
    main()