"""
Stem Separation Module
Uses Demucs for audio source separation
"""

import subprocess
from pathlib import Path
import sys


def separate_stems(input_file, output_dir, two_stem=True):
    """
    Separate audio into stems using Demucs (CPU mode)
    
    Args:
        input_file: Path to input audio file
        output_dir: Directory to save stems
        two_stem: If True, only separate vocals/instrumental (faster)
                  If False, separate into 6 stems (slower)
    
    Returns:
        Path to directory containing stems
    """
    
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    if two_stem:
        print(f"   Using 2-stem mode (vocal + instrumental)")
        print(f"   Estimated time: 2-4 minutes")
        
        cmd = [
            'demucs',
            '--two-stems=vocals',
            '-o', str(output_path),
            '--mp3',
            '--mp3-bitrate=320',
            '--jobs=0',  # Use all CPU cores
            input_file
        ]
    else:
        print(f"   Using 6-stem mode (vocal, drums, bass, guitar, keys, other)")
        print(f"   Estimated time: 8-12 minutes")
        
        cmd = [
            'demucs',
            '-n', 'htdemucs_6s',
            '-o', str(output_path),
            '--mp3',
            '--mp3-bitrate=320',
            '--jobs=0',
            input_file
        ]
    
    try:
        # Run Demucs with live output
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            bufsize=1
        )
        
        # Print progress
        for line in process.stdout:
            line = line.strip()
            if line:
                # Filter for progress info
                if '%' in line or 'Separated' in line or 'Processing' in line:
                    print(f"   {line}")
        
        process.wait()
        
        if process.returncode != 0:
            raise Exception(f"Demucs exited with code {process.returncode}")
        
        # Locate output directory
        # Demucs structure: output_dir/model_name/track_name/stems
        song_stem = Path(input_file).stem
        
        # Try different possible structures
        possible_paths = [
            output_path / 'htdemucs' / song_stem,
            output_path / 'htdemucs_6s' / song_stem,
            output_path / song_stem
        ]
        
        for stems_dir in possible_paths:
            if stems_dir.exists():
                # Verify stems are present
                stem_files = list(stems_dir.glob('*.mp3'))
                if stem_files:
                    return stems_dir
        
        # If we get here, something went wrong
        raise Exception(f"Could not find output stems in {output_path}")
        
    except FileNotFoundError:
        print("\n[FAIL] Error: 'demucs' command not found")
        print("   Make sure Demucs is installed:")
        print("   pip install demucs")
        sys.exit(1)
        
    except Exception as e:
        raise Exception(f"Stem separation failed: {str(e)}")


def list_stems(stems_dir):
    """List all stem files in directory"""
    stems_dir = Path(stems_dir)
    return sorted(stems_dir.glob('*.mp3'))


if __name__ == "__main__":
    # Test module
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python stems.py <audio_file>")
        sys.exit(1)
    
    test_file = sys.argv[1]
    test_output = "./test_output"
    
    print(f"Testing stem separation on: {test_file}")
    result = separate_stems(test_file, test_output, two_stem=True)
    
    print(f"\nStems created in: {result}")
    print("\nFiles:")
    for stem in list_stems(result):
        print(f"  • {stem.name}")