import json
import os
import asyncio
import time
import traceback
from ahserver.filestorage import FileStorage
from urllib.parse import quote

try:
    DEMUCS_WRAPPER = '/tmp/demucs_wrapper.py'
    PYTHON_BIN = '/share/vllm-0.8.5/bin/python'
    OUTPUT_BASE_DIR = '/tmp/demucs_output'
    
    os.makedirs(OUTPUT_BASE_DIR, exist_ok=True)
    
    fs = FileStorage()
    filesroot = os.path.abspath(fs.root)  # e.g. /tmp
    
    input_file = None
    web_path = None
    
    # 方式1: 文件上传 (multipart/form-data)
    if 'audio_file' in params_kw:
        web_path = params_kw['audio_file']
        if isinstance(web_path, str) and len(web_path) > 0:
            input_file = fs.realPath(web_path)
    
    # 方式2: 文件路径
    elif 'filepath' in params_kw or 'audio_path' in params_kw:
        input_file = params_kw.get('filepath') or params_kw.get('audio_path')
    
    if not input_file:
        return json.dumps({
            'status': 'error', 
            'message': 'Missing parameter: audio_file (upload) or filepath (path)'
        }, ensure_ascii=False)
    
    if not os.path.isfile(input_file):
        return json.dumps({
            'status': 'error', 
            'message': f'File not found: {input_file}',
            'web_path': web_path
        }, ensure_ascii=False)
    
    run_id = f'demucs_{int(time.time() * 1000)}'
    output_dir = os.path.join(OUTPUT_BASE_DIR, run_id)
    os.makedirs(output_dir, exist_ok=True)

    cmd = [PYTHON_BIN, DEMUCS_WRAPPER, '--two-stems=vocals', '-o', output_dir, input_file]

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

    if proc.returncode != 0:
        error_msg = stderr.decode('utf-8', errors='replace').strip()
        return json.dumps({
            'status': 'error', 
            'message': f'Demucs failed: {error_msg}'
        }, ensure_ascii=False)
    
    filename_no_ext = os.path.splitext(os.path.basename(input_file))[0]
    stems_dir = os.path.join(output_dir, 'htdemucs', filename_no_ext)
    vocals_path = os.path.join(stems_dir, 'vocals.wav')
    no_vocals_path = os.path.join(stems_dir, 'no_vocals.wav')

    if not os.path.isfile(vocals_path):
        return json.dumps({
            'status': 'error', 
            'message': f'Output not found: {vocals_path}'
        }, ensure_ascii=False)
    
    # 计算相对于 filesroot 的 web path
    vocals_webpath = fs.webpath(vocals_path)
    no_vocals_webpath = fs.webpath(no_vocals_path)
    
    # 构建下载 URL (通过 idfile 端点)
    vocals_url = f'/idfile?path={quote(vocals_webpath)}'
    no_vocals_url = f'/idfile?path={quote(no_vocals_webpath)}'
    
    return json.dumps({
        'status': 'success',
        'vocals_url': vocals_url,
        'no_vocals_url': no_vocals_url,
        'input_file': input_file
    }, ensure_ascii=False)

except Exception as e:
    return json.dumps({
        'status': 'error', 
        'message': str(e), 
        'traceback': traceback.format_exc()
    }, ensure_ascii=False)