import os
import sys
import subprocess
import torch
from pathlib import Path

# Project paths
PROJECT_ROOT = Path(__file__).resolve().parent
VENV_PYTHON = PROJECT_ROOT / "venv" / "Scripts" / "python.exe"

def print_banner():
    banner = """
  +----------------------------------------------------------+
  |         *** TITANAUDIO PRO : THE PHANTOM ENGINE ***      |
  |            -- GPU-Accelerated Audio AI Pipeline --       |
  +----------------------------------------------------------+
    """
    print(banner)

def print_sys_info():
    print("  [SYSTEM STATUS]")
    cuda_avail = torch.cuda.is_available()
    print(f"  |-- PyTorch Version : {torch.__version__}")
    print(f"  |-- CUDA Available  : {'YES (GPU)' if cuda_avail else 'NO (CPU)'}")
    if cuda_avail:
        device_name = torch.cuda.get_device_name(0)
        total_mem = round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2)
        print(f"  |-- GPU Device      : {device_name}")
        print(f"  +-- Total VRAM      : {total_mem} GB")
    print("  " + "-"*58 + "\n")

def run_stage(stage_num, stage_name, script_path, args=[]):
    print(f"  => [STEP {stage_num}/3] {stage_name}...")
    cmd = [str(VENV_PYTHON), str(script_path)] + args
    print(f"  |   Executing: {script_path.name}")
    
    # Run the stage script
    result = subprocess.run(cmd, capture_output=False)
    
    if result.returncode == 0:
        print(f"  +==> [SUCCESS] {stage_name} completed successfully!\n")
        return True
    else:
        print(f"  +==> [ERROR] {stage_name} failed with exit code {result.returncode}\n")
        return False

def main():
    print_banner()
    print_sys_info()

    if len(sys.argv) < 2:
        print("  [ERROR] No input file provided!")
        print("  Usage: python main_pipeline.py <input_audio_file>\n")
        return

    input_file = sys.argv[1]
    input_path = Path(input_file).resolve()
    
    if not input_path.exists():
        print(f"  [ERROR] Input file not found: {input_file}\n")
        return

    # --- STAGE 2: Separation ---
    sep_script = PROJECT_ROOT / "2_Stem_Separation" / "separate_gpu.py"
    if not run_stage(1, "Deep AI Stem Separation (Demucs)", sep_script, [str(input_path)]):
        print("  [CRITICAL] Separation failed. Pipeline aborted.\n")
        return

    # --- STAGE 3: Transformation ---
    fx_script = PROJECT_ROOT / "3_Transformation" / "apply_fx_gpu.py"
    if not run_stage(2, "Multi-Layer FX Transformation (Pedalboard)", fx_script):
        print("  [CRITICAL] FX transformation failed. Pipeline aborted.\n")
        return

    # --- STAGE 5 & 6: Mix & Master ---
    master_script = PROJECT_ROOT / "5_Mastering" / "master_final.py"
    if not run_stage(3, "High-Fidelity AI Mastering (Matchering)", master_script):
        print("  [CRITICAL] AI Mastering failed. Pipeline aborted.\n")
        return

    # Final Success Summary
    print("  " + "="*58)
    print("  *** TITAN ENGINE SUCCESS: Audio processed flawlessly! ***")
    print(f"  * Final Mix: Outputs/Final_Mix.wav")
    print(f"  * Mastered : Outputs/Titan_Master_Final.wav")
    print("  " + "="*58 + "\n")

if __name__ == "__main__":
    main()