import os
import json
import asyncio
import time
from appPublic.log import debug, error

async def run_separate(task_obj, payload):
    """
    Run demucs vocal/accompaniment separation.

    payload:
        audio_path (str, required): Path to input audio file
        output_dir (str, optional): Output directory, default /tmp/demucs_{task_id}
    """
    audio_path = payload.get('audio_path')
    if not audio_path:
        raise ValueError('audio_path is required')

    if not os.path.isfile(audio_path):
        raise FileNotFoundError(f'Audio file not found: {audio_path}')

    task_id = payload.get('task_id', str(int(time.time())))
    output_dir = payload.get('output_dir', f'/tmp/demucs_{task_id}')

    gpu_id = task_obj.gpu_id
    basename = os.path.splitext(os.path.basename(audio_path))[0]

    # Expected output paths from demucs
    result_dir = os.path.join(output_dir, 'htdemucs', basename)
    vocals_path = os.path.join(result_dir, 'vocals.wav')
    no_vocals_path = os.path.join(result_dir, 'no_vocals.wav')

    # Build the command
    env = os.environ.copy()
    env['CUDA_VISIBLE_DEVICES'] = str(gpu_id)

    cmd = [
        '/data/ymq/demucs_venv/bin/python', '-m', 'demucs',
        '--two-stems', 'vocals',
        audio_path,
        '-o', output_dir
    ]

    debug(f'[demucs] Running separation: audio={audio_path}, output={output_dir}, gpu={gpu_id}')
    debug(f'[demucs] Command: {" ".join(cmd)}')

    start_time = time.time()

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env=env
    )

    stdout, stderr = await proc.communicate()
    elapsed = time.time() - start_time

    if proc.returncode != 0:
        stderr_text = stderr.decode('utf-8', errors='replace')
        stdout_text = stdout.decode('utf-8', errors='replace')
        error(f'[demucs] Process failed (rc={proc.returncode})')
        error(f'[demucs] stdout: {stdout_text[-2000:]}')
        error(f'[demucs] stderr: {stderr_text[-2000:]}')
        raise RuntimeError(
            f'Demucs separation failed (rc={proc.returncode}): {stderr_text[-500:]}'
        )

    # Verify output files exist
    if not os.path.isfile(vocals_path):
        raise FileNotFoundError(f'Expected vocals output not found: {vocals_path}')
    if not os.path.isfile(no_vocals_path):
        raise FileNotFoundError(f'Expected no_vocals output not found: {no_vocals_path}')

    vocals_size = os.path.getsize(vocals_path)
    no_vocals_size = os.path.getsize(no_vocals_path)

    debug(f'[demucs] Separation complete in {elapsed:.1f}s')
    debug(f'[demucs] vocals.wav: {vocals_size} bytes, no_vocals.wav: {no_vocals_size} bytes')

    return {
        'vocals_path': vocals_path,
        'no_vocals_path': no_vocals_path,
        'duration': round(elapsed, 2),
        'output_dir': output_dir,
        'model': 'htdemucs'
    }