name: CI

on:
  push:
    branches: [main, master]
    paths:
      - '**/*.cpp'
      - '**/*.h'
      - '**/*.hpp'
      - '**/CMakeLists.txt'
      - '.github/workflows/**'
  pull_request:
    types: [opened, synchronize, reopened]
  workflow_dispatch:

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
  cancel-in-progress: true

env:
  # HuggingFace model info
  HF_MODEL_REPO: GaboxR67/MelBandRoformers
  HF_CHECKPOINT_PATH: melbandroformers/vocals/voc_fv6.ckpt
  HF_CONFIG_PATH: melbandroformers/vocals/voc_gabox.yaml
  # Music-Source-Separation-Training repo
  MSST_REPO: https://github.com/ZFTurbo/Music-Source-Separation-Training.git
  # Enable sccache GitHub Actions cache
  SCCACHE_GHA_ENABLED: "true"

jobs:
  # ===========================================================================
  # Prepare: Generate test data (runs once, shared via artifacts)
  # ===========================================================================
  prepare-test-data:
    runs-on: ubuntu-latest
    
    steps:
      - name: Checkout
        uses: actions/checkout@v4
      
      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'
          
      - name: Clone MSST Repository
        run: git clone --depth 1 ${{ env.MSST_REPO }} msst
        
      - name: Install Dependencies
        run: |
          pip install torch torchaudio --index-url https://download.pytorch.org/whl/cpu
          pip install huggingface_hub scipy soundfile gguf librosa ml_collections einops pyyaml numpy tqdm beartype rotary_embedding_torch
          
      - name: Download Model from HuggingFace
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
        run: |
          python -c "
          from huggingface_hub import hf_hub_download
          import os
          token = os.environ.get('HF_TOKEN') or None
          hf_hub_download('${{ env.HF_MODEL_REPO }}', '${{ env.HF_CHECKPOINT_PATH }}', 
                          local_dir='./model', token=token)
          hf_hub_download('${{ env.HF_MODEL_REPO }}', '${{ env.HF_CONFIG_PATH }}',
                          local_dir='./model', token=token)
          "
          
      - name: Generate Test Audio
        run: |
          python scripts/generate_test_audio.py --output test_audio.wav --duration 5.0 --sample-rate 44100
          
      - name: Generate Test Data
        run: |
          python scripts/generate_test_data.py \
            --model-repo msst \
            --audio test_audio.wav \
            --checkpoint model/${{ env.HF_CHECKPOINT_PATH }} \
            --config model/${{ env.HF_CONFIG_PATH }} \
            --output test_data
            
      - name: Convert Model to GGUF
        run: |
          python scripts/convert_to_gguf.py \
            --ckpt model/${{ env.HF_CHECKPOINT_PATH }} \
            --config model/${{ env.HF_CONFIG_PATH }} \
            --out model.gguf \
            --dtype fp32
            
      - name: Upload Test Data Artifact
        uses: actions/upload-artifact@v4
        with:
          name: test-data
          path: |
            test_data/
            model.gguf
            test_audio.wav
          retention-days: 1

  # ===========================================================================
  # Build Matrix: Core Platforms + Vulkan
  # ===========================================================================
  build:
    needs: prepare-test-data
    strategy:
      fail-fast: false
      matrix:
        include:
          # Tier 1: Core Platforms (CPU)
          - { name: linux-x64-cpu,    os: ubuntu-22.04,     backend: cpu,    test: true  }
          - { name: linux-arm64-cpu,  os: ubuntu-22.04-arm, backend: cpu,    test: true  }
          - { name: macos-arm64,      os: macos-latest,     backend: cpu,    test: true  }
          - { name: macos-x64,        os: macos-15-intel,   backend: cpu,    test: true  }
          - { name: windows-x64-msvc, os: windows-2025,     backend: cpu,    test: true  }
          # Tier 2: Vulkan Backend
          - { name: linux-vulkan,     os: ubuntu-24.04,     backend: vulkan, test: true  }
          - { name: windows-vulkan,   os: windows-2025,     backend: vulkan, test: true  }
          
    runs-on: ${{ matrix.os }}
    
    steps:
      - name: Checkout
        uses: actions/checkout@v4
      
      - name: Clone GGML
        run: git clone --depth 1 https://github.com/ggerganov/ggml.git ggml
      
      - name: Download Test Data
        uses: actions/download-artifact@v4
        with:
          name: test-data
          
      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'
          
      - name: Setup MSVC
        if: runner.os == 'Windows'
        uses: ilammy/msvc-dev-cmd@v1

      - name: Install Python Dependencies
        run: pip install numpy scipy
      
      - name: Setup sccache
        uses: mozilla-actions/sccache-action@v0.0.9
        with:
          version: "v0.12.0"
          
      # ----- Linux Dependencies -----
      - name: Install Dependencies (Linux)
        if: runner.os == 'Linux'
        run: |
          sudo apt-get update
          sudo apt-get install -y build-essential cmake
          
      - name: Install Vulkan SDK (Linux)
        if: matrix.backend == 'vulkan' && runner.os == 'Linux'
        run: |
          sudo apt-get install -y libvulkan-dev glslc mesa-vulkan-drivers
          
      # ----- macOS Dependencies -----  
      - name: Install Dependencies (macOS)
        if: runner.os == 'macOS'
        run: brew install cmake
        
      # ----- Windows Dependencies -----
      - name: Install Dependencies (Windows)
        if: runner.os == 'Windows'
        run: choco install ninja -y
        
      - name: Install Vulkan SDK (Windows)
        if: matrix.backend == 'vulkan' && runner.os == 'Windows'
        run: |
          $VK_VERSION = "1.4.313.2"
          curl.exe -o VulkanSDK.exe -L "https://sdk.lunarg.com/sdk/download/${VK_VERSION}/windows/vulkansdk-windows-X64-${VK_VERSION}.exe"
          Start-Process -FilePath .\VulkanSDK.exe -ArgumentList "--accept-licenses --default-answer --confirm-command install" -Wait
          Add-Content $env:GITHUB_ENV "VULKAN_SDK=C:\VulkanSDK\${VK_VERSION}"
          Add-Content $env:GITHUB_PATH "C:\VulkanSDK\${VK_VERSION}\bin"
          
      # ----- Configure -----
      - name: Configure (Unix)
        if: runner.os != 'Windows'
        run: |
          cmake -B build \
            -DCMAKE_BUILD_TYPE=Release \
            -DGGML_DIR=ggml \
            -DGGML_CUDA=OFF \
            -DGGML_VULKAN=${{ matrix.backend == 'vulkan' && 'ON' || 'OFF' }} \
            -DCMAKE_C_COMPILER_LAUNCHER=sccache \
            -DCMAKE_CXX_COMPILER_LAUNCHER=sccache \
            -DMBR_BUILD_TESTS=ON \
            -DMBR_BUILD_CLI=ON
            
      - name: Configure (Windows)
        if: runner.os == 'Windows'
        run: |
          cmake -B build -G "Ninja Multi-Config" `
            -DGGML_DIR=ggml `
            -DGGML_CUDA=OFF `
            -DGGML_VULKAN=${{ matrix.backend == 'vulkan' && 'ON' || 'OFF' }} `
            -DCMAKE_C_COMPILER_LAUNCHER=sccache `
            -DCMAKE_CXX_COMPILER_LAUNCHER=sccache `
            -DMBR_BUILD_TESTS=ON `
            -DMBR_BUILD_CLI=ON
            
      # ----- Build -----
      - name: Build (Unix)
        if: runner.os != 'Windows'
        run: cmake --build build --config Release -j $(nproc 2>/dev/null || sysctl -n hw.logicalcpu)
        
      - name: Build (Windows)
        if: runner.os == 'Windows'
        run: cmake --build build --config Release -j $env:NUMBER_OF_PROCESSORS
        
      # ----- Unit Tests -----
      - name: Run Unit Tests
        if: matrix.test
        env:
          MBR_MODEL_PATH: ${{ github.workspace }}/model.gguf
          MBR_TEST_DATA_DIR: ${{ github.workspace }}/test_data
          MBR_FORCE_CPU: ${{ runner.os == 'macOS' && '1' || '' }}
        run: ctest --test-dir build -C Release -V --output-on-failure --timeout 300
        
      # ----- CLI Tests -----
      - name: Test CLI
        if: matrix.test
        shell: bash
        env:
          MBR_MODEL_PATH: ${{ github.workspace }}/model.gguf
          MBR_FORCE_CPU: ${{ runner.os == 'macOS' && '1' || '' }}
        run: |
          echo "=== CLI Test Suite ==="
          
          # Dynamically find CLI executable
          CLI_NAME="mel_band_roformer-cli"
          if [[ "$RUNNER_OS" == "Windows" ]]; then CLI_NAME="mel_band_roformer-cli.exe"; fi
          
          echo "Searching for $CLI_NAME..."
          CLI_PATH=$(find build -name "$CLI_NAME" | head -n 1)
          
          if [[ -z "$CLI_PATH" ]]; then
            echo "Error: CLI executable not found!"
            find build
            exit 1
          fi
          
          echo "Found CLI at: $CLI_PATH"
          chmod +x "$CLI_PATH"
          
          # Setup execution variables
          CLI_DIR=$(dirname "$CLI_PATH")
          CLI_EXE="./$CLI_NAME"
          
          # Debug Info
          echo "Debug: Listing directory $CLI_DIR"
          ls -l "$CLI_DIR"
          
          echo "Debug: Checking dependencies"
          if command -v ldd >/dev/null; then
             ldd "$CLI_PATH" || echo "ldd returned error"
          fi
          
          # Define runner helper
          run_cli() {
              echo "Executing: ./$CLI_NAME $@"
              (cd "$CLI_DIR" && ./$CLI_NAME "$@")
          }
          
          # Debug PATH
          echo "PATH: $PATH"
          
          # 1. Test --help
          echo "[1/4] Testing --help..."
          run_cli --help
          
          # 2. Test with missing arguments (should fail)
          echo "[2/4] Testing error handling..."
          if run_cli 2>/dev/null; then
            echo "ERROR: CLI should fail without arguments"
            exit 1
          fi
          
          # 3. Use existing test audio
          echo "[3/4] Using generated test audio..."
          cp test_audio.wav cli_test_input.wav
          
          # 4. Run full inference
          echo "[4/4] Running inference..."
          # Use absolute paths for input/output to avoid directory issues
          ABS_MODEL=$(readlink -f "$MBR_MODEL_PATH" || echo "$(pwd)/model.gguf")
          ABS_INPUT=$(readlink -f cli_test_input.wav || echo "$(pwd)/cli_test_input.wav")
          ABS_OUTPUT=$(readlink -f cli_test_output.wav || echo "$(pwd)/cli_test_output.wav")
          
          # Convert paths for Windows MSVC executables
          if [[ "$RUNNER_OS" == "Windows" ]]; then
              ABS_MODEL=$(cygpath -w "$ABS_MODEL")
              ABS_INPUT=$(cygpath -w "$ABS_INPUT")
              ABS_OUTPUT=$(cygpath -w "$ABS_OUTPUT")
          fi
          
          echo "Running with model: $ABS_MODEL"
          run_cli "$ABS_MODEL" "$ABS_INPUT" "$ABS_OUTPUT" --chunk-size 88200 --overlap 2
          
          # Verify output exists and has reasonable size
          if [[ ! -f cli_test_output.wav ]]; then
            echo "ERROR: Output file not created"
            exit 1
          fi
          
          OUTPUT_SIZE=$(stat -c%s cli_test_output.wav 2>/dev/null || stat -f%z cli_test_output.wav)
          if [[ $OUTPUT_SIZE -lt 1000 ]]; then
            echo "ERROR: Output file too small: $OUTPUT_SIZE bytes"
            exit 1
          fi
          
          echo "=== CLI Tests Passed ==="
          
      # ----- Upload Artifacts -----
      - name: Upload Build Artifacts
        uses: actions/upload-artifact@v4
        with:
          name: build-${{ matrix.name }}
          path: |
            build/bin/
            build/lib*/
            build/*.dll
            build/*.so
            build/*.dylib
            build/mel_band_roformer-cli*
            build/Release/
          retention-days: 7
      
      # ----- Prepare Release Artifact -----
      - name: Prepare Release Artifact (Unix)
        if: runner.os != 'Windows'
        shell: bash
        run: |
          # Create release directory
          mkdir -p release/mel-band-roformer
          
          # Find and copy CLI executable
          CLI_PATH=$(find build -name "mel_band_roformer-cli" -type f | head -n 1)
          if [[ -n "$CLI_PATH" ]]; then
            cp "$CLI_PATH" release/mel-band-roformer/
            chmod +x release/mel-band-roformer/mel_band_roformer-cli
          fi
          
          # Copy shared libraries if exist (only real files, not symlinks)
          find build \( -name "*.so*" -o -name "*.dylib" \) -type f ! -type l | while read lib; do
            cp "$lib" release/mel-band-roformer/ 2>/dev/null || true
          done
          
          # List contents
          echo "Release artifact contents:"
          ls -lh release/mel-band-roformer/
      
      - name: Prepare Release Artifact (Windows)
        if: runner.os == 'Windows'
        shell: pwsh
        run: |
          # Create release directory
          New-Item -ItemType Directory -Force -Path "release\mel-band-roformer"
          
          # Find and copy CLI executable
          $CliPath = Get-ChildItem -Path build -Filter "mel_band_roformer-cli.exe" -Recurse -File | Select-Object -First 1
          if ($CliPath) {
            Copy-Item $CliPath.FullName "release\mel-band-roformer\"
          }
          
          # Copy DLL files
          Get-ChildItem -Path build -Filter "*.dll" -Recurse -File | ForEach-Object {
            Copy-Item $_.FullName "release\mel-band-roformer\" -ErrorAction SilentlyContinue
          }
          
          # List contents
          Write-Host "Release artifact contents:"
          Get-ChildItem "release\mel-band-roformer" | Format-Table Name, Length
      
      - name: Upload Release Artifact
        uses: actions/upload-artifact@v4
        with:
          name: MelBandRoformer-${{ matrix.name }}
          path: release/mel-band-roformer/
          retention-days: 30

  # ===========================================================================
  # CUDA Build: Linux (Compile Verification Only)
  # ===========================================================================
  build-cuda-linux:
    name: build-cuda-linux-${{ matrix.cuda_version }}
    runs-on: ${{ matrix.os }}
    strategy:
      fail-fast: false
      matrix:
        include:
          - { cuda_version: "11.8.0", os: ubuntu-22.04 }
          - { cuda_version: "12.9.1", os: ubuntu-latest }
          - { cuda_version: "13.1.0", os: ubuntu-latest }
    
    steps:  
      - name: Checkout
        uses: actions/checkout@v4
      
      - name: Clone GGML
        run: git clone --depth 1 https://github.com/ggerganov/ggml.git ggml

      - name: Install CUDA Toolkit
        uses: Jimver/cuda-toolkit@master
        with:
          cuda: ${{ matrix.cuda_version }}
          method: network
          sub-packages: '["nvcc", "cudart", "thrust"]'
          non-cuda-sub-packages: '["libcublas", "libcublas-dev"]'
      
      - name: Install Dependencies
        run: |
          sudo apt-get install -y cmake build-essential ninja-build

      - name: Setup sccache
        uses: mozilla-actions/sccache-action@v0.0.9
          
      - name: Configure
        run: |
          ls -ld ggml
          # Minimal architectures for compile verification
          # Just verify compilation works, not for distribution
          CUDA_ARCHS="75;86"
          echo "Verifying build for CUDA architectures: $CUDA_ARCHS"
          cmake -B build -G Ninja \
            -DCMAKE_BUILD_TYPE=Release \
            -DGGML_DIR=ggml \
            -DGGML_CUDA=ON \
            -DGGML_CUDA_FORCE_MMQ=ON \
            -DCMAKE_CUDA_ARCHITECTURES="$CUDA_ARCHS" \
            -DCMAKE_C_COMPILER_LAUNCHER=sccache \
            -DCMAKE_CXX_COMPILER_LAUNCHER=sccache \
            -DCMAKE_CUDA_COMPILER_LAUNCHER=sccache \
            -DMBR_BUILD_TESTS=OFF \
            -DMBR_BUILD_CLI=ON
            
      - name: Build
        run: cmake --build build --config Release -j $(nproc)
        
      - name: Upload Artifacts
        uses: actions/upload-artifact@v4
        with:
          name: build-linux-cuda-${{ matrix.cuda_version }}
          path: |
            build/bin/
            build/lib*/
            build/*.so
          retention-days: 7
      
      # ----- Prepare Release Artifact -----
      - name: Prepare Release Artifact
        run: |
          # Create release directory
          mkdir -p release/mel-band-roformer
          
          # Find and copy CLI executable
          CLI_PATH=$(find build -name "mel_band_roformer-cli" -type f | head -n 1)
          if [[ -n "$CLI_PATH" ]]; then
            cp "$CLI_PATH" release/mel-band-roformer/
            chmod +x release/mel-band-roformer/mel_band_roformer-cli
          fi
          
          # Copy shared libraries
          find build -name "*.so*" | while read lib; do
            cp "$lib" release/mel-band-roformer/ 2>/dev/null || true
          done
          
          # List contents
          echo "Release artifact contents:"
          ls -lh release/mel-band-roformer/
      
      - name: Upload Release Artifact
        uses: actions/upload-artifact@v4
        with:
          name: MelBandRoformer-linux-cuda-${{ matrix.cuda_version }}
          path: release/mel-band-roformer/
          retention-days: 30

  # ===========================================================================
  # CUDA Build: Windows (Compile Only - No GPU for testing)
  # ===========================================================================
  build-cuda-windows:
    name: build-cuda-windows-${{ matrix.cuda_version }}
    runs-on: windows-2022
    strategy:
      fail-fast: false
      matrix:
        cuda_version: ["11.8.0", "12.9.1", "13.1.0"]
    
    env:
      CUDA_VERSION: ${{ matrix.cuda_version }}
    
    steps:
      - name: Checkout
        uses: actions/checkout@v4
      
      - name: Setup MSVC
        if: runner.os == 'Windows'
        uses: ilammy/msvc-dev-cmd@v1

      - name: Clone GGML
        run: git clone --depth 1 https://github.com/ggerganov/ggml.git ggml
      
      - name: Install CUDA Toolkit
        if: ${{ matrix.cuda_version != '13.1.0' }}
        uses: Jimver/cuda-toolkit@master
        with:
          cuda: ${{ matrix.cuda_version }}
          method: network
          sub-packages: '["nvcc", "cudart", "cublas", "cublas_dev", "thrust", "visual_studio_integration"]'
        
      - name: Install CUDA Toolkit(13.1.0)
        if: ${{ matrix.cuda_version == '13.1.0' }}
        uses: Jimver/cuda-toolkit@master
        with:
          cuda: ${{ matrix.cuda_version }}
          method: network
          sub-packages: '["nvcc", "cudart", "cublas", "cublas_dev", "nvrtc", "nvrtc_dev", "crt", "nvvm", "visual_studio_integration"]'

      - name: Install Ninja
        run: choco install ninja -y

      - name: Setup sccache
        uses: mozilla-actions/sccache-action@v0.0.9
        
      - name: Configure and Build
        shell: pwsh
        run: |
          # Consumer GPU architectures:
          # 61=Pascal (GTX 10), 75=Turing (RTX 20/GTX 16), 86=Ampere (RTX 30), 89=Ada (RTX 40), 120=Blackwell (RTX 50)
          $cudaVersion = "${{ matrix.cuda_version }}"
          if ($cudaVersion -match "^11\.") {
            # CUDA 11.x doesn't support arch 120
            $cudaArchs = "61;75;86;89"
            $env:CUDAFLAGS = "-allow-unsupported-compiler -D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH -D_DISABLE_CONSTEXPR_MUTEX_CONSTRUCTOR"
          } else {
            # CUDA 12+: Add RTX 50 (Blackwell)
            $cudaArchs = "61;75;86;89;120"
            $env:CUDAFLAGS = ""
          }
          Write-Host "Building for CUDA architectures: $cudaArchs"
          
          cmake -B build -G "Ninja Multi-Config" `
            -DGGML_DIR=ggml `
            -DGGML_CUDA=ON `
            -DGGML_CUDA_FORCE_MMQ=ON `
            "-DCMAKE_CUDA_ARCHITECTURES=$cudaArchs" `
            -DCMAKE_C_COMPILER_LAUNCHER=sccache `
            -DCMAKE_CXX_COMPILER_LAUNCHER=sccache `
            -DMBR_BUILD_TESTS=OFF `
            -DMBR_BUILD_CLI=ON
          cmake --build build --config Release -j $env:NUMBER_OF_PROCESSORS
        
      - name: Upload Artifacts
        uses: actions/upload-artifact@v4
        with:
          name: build-windows-cuda-${{ matrix.cuda_version }}
          path: |
            build/bin/
            build/Release/
            build/*.dll
          retention-days: 7
      
      # ----- Prepare Release Artifact -----
      - name: Prepare Release Artifact
        shell: pwsh
        run: |
          # Create release directory
          New-Item -ItemType Directory -Force -Path "release\mel-band-roformer"
          
          # Find and copy CLI executable
          $CliPath = Get-ChildItem -Path build -Filter "mel_band_roformer-cli.exe" -Recurse -File | Select-Object -First 1
          if ($CliPath) {
            Copy-Item $CliPath.FullName "release\mel-band-roformer\"
          }
          
          # Copy DLL files
          Get-ChildItem -Path build -Filter "*.dll" -Recurse -File | ForEach-Object {
            Copy-Item $_.FullName "release\mel-band-roformer\" -ErrorAction SilentlyContinue
          }
          
          # List contents
          Write-Host "Release artifact contents:"
          Get-ChildItem "release\mel-band-roformer" | Format-Table Name, Length
      
      - name: Upload Release Artifact
        uses: actions/upload-artifact@v4
        with:
          name: MelBandRoformer-windows-cuda-${{ matrix.cuda_version }}
          path: release\mel-band-roformer\
          retention-days: 30