import os
import sys
import requests
from gradio_client import Client, handle_file

raw_prompt = os.environ.get('PROMPT', '')
run_id = os.environ.get('RUN_ID', '')
space_url = os.environ.get('SPACE_URL', '')
github_run_id = os.environ.get('GITHUB_RUN_ID', '')

def report_failure(error_msg):
    try:
        requests.post(
            f"{space_url}/api/webhook/fail",
            json={
                "run_id": run_id,
                "error": error_msg,
                "event_type": "separate-audio",
                "client_payload": {
                    "prompt": raw_prompt,
                    "run_id": run_id,
                    "space_url": space_url
                },
                "github_run_id": github_run_id
            },
            timeout=15
        )
    except Exception as e:
        print(f"Failed to report failure: {e}")

print('1. Decoding configuration from separate payload...')
if not raw_prompt.startswith("VOICECONFIG_SEPARATE_"):
    err_str = "Error: Invalid separate configuration payload signature."
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

config_str = raw_prompt[len("VOICECONFIG_SEPARATE_"):]
parts = config_str.split("_")
config = {}
i = 0
while i < len(parts) - 1:
    key = parts[i]
    val = parts[i+1]
    config[key] = val
    i += 2

user_run_id = config.get("userRunId", run_id)
ext = config.get("ext", "mp3")
stem = config.get("stem", "vocal")

# تبدیل مقادیر رشته‌ای به مقادیر و متغیرهای اصلی پایتون
main = config.get("main", "false").lower() == "true"
dereverb = config.get("dereverb", "false").lower() == "true"
vocal_effects = config.get("vocEff", "false").lower() == "true"
background_effects = config.get("bgEff", "false").lower() == "true"

vocal_reverb_room_size = float(config.get("vRevRoom", "0.15"))
vocal_reverb_damping = float(config.get("vRevDamp", "0.7"))
vocal_reverb_dryness = float(config.get("vRevDry", "0.8"))
vocal_reverb_wet_level = float(config.get("vRevWet", "0.2"))
vocal_delay_seconds = float(config.get("vDelaySec", "0.0"))
vocal_delay_mix = float(config.get("vDelayMix", "0.0"))
vocal_compressor_threshold_db = float(config.get("vCompThresh", "-15"))
vocal_compressor_ratio = float(config.get("vCompRatio", "4"))
vocal_compressor_attack_ms = float(config.get("vCompAttack", "1"))
vocal_compressor_release_ms = float(config.get("vCompRelease", "100"))
vocal_gain_db = float(config.get("vGain", "0"))

background_highpass_freq = float(config.get("bgHigh", "120"))
background_lowpass_freq = float(config.get("bgLow", "11000"))
background_reverb_room_size = float(config.get("bgRevRoom", "0.1"))
background_reverb_damping = float(config.get("bgRevDamp", "0.5"))
background_reverb_wet_level = float(config.get("bgRevWet", "0.25"))
background_compressor_threshold_db = float(config.get("bgCompThresh", "-15"))
background_compressor_ratio = float(config.get("bgCompRatio", "4"))
background_compressor_attack_ms = float(config.get("bgCompAttack", "15"))
background_compressor_release_ms = float(config.get("bgCompRelease", "60"))
background_gain_db = float(config.get("bgGain", "0"))
target_format = config.get("format", "WAV")

input_audio_url = f"{space_url}/static/images/{user_run_id}_input.{ext}"
local_input = f"input.{ext}"

print("2. Downloading source audio from host...")
try:
    r_input = requests.get(input_audio_url, timeout=60)
    if r_input.status_code != 200:
        raise Exception(f"Input audio download failed. Status: {r_input.status_code}")
    with open(local_input, 'wb') as f:
        f.write(r_input.content)
except Exception as download_err:
    err_str = f"Error downloading source files: {download_err}"
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

print("3. Connecting to Audio Separator Space...")
try:
    client = Client("https://r3gm-audio-separator.hf.space/")
    
    print("4. Executing separation engine...")
    result = client.predict(
        handle_file(local_input),
        [stem],
        main,
        dereverb,
        vocal_effects,
        background_effects,
        vocal_reverb_room_size,
        vocal_reverb_damping,
        vocal_reverb_dryness,
        vocal_reverb_wet_level,
        vocal_delay_seconds,
        vocal_delay_mix,
        vocal_compressor_threshold_db,
        vocal_compressor_ratio,
        vocal_compressor_attack_ms,
        vocal_compressor_release_ms,
        vocal_gain_db,
        background_highpass_freq,
        background_lowpass_freq,
        background_reverb_room_size,
        background_reverb_damping,
        background_reverb_wet_level,
        background_compressor_threshold_db,
        background_compressor_ratio,
        background_compressor_attack_ms,
        background_compressor_release_ms,
        background_gain_db,
        target_format,
        fn_index=3
    )

    def parse_file_response(f):
        if not f: return None
        if isinstance(f, (list, tuple)):
            if len(f) > 0:
                return parse_file_response(f[0])
        if isinstance(f, dict):
            return f.get('path') or f.get('name')
        return str(f)

    final_audio_path = parse_file_response(result)

    if not final_audio_path or not os.path.exists(str(final_audio_path)):
        raise Exception("Audio separation output file was not found or invalid.")

    print("5. Uploading result file back...")
    ext_out = target_format.lower()
    with open(final_audio_path, 'rb') as f:
        res_upload = requests.post(
            f'{space_url}/api/webhook/upload', 
            data={'run_id': run_id, 'github_run_id': github_run_id, 'ext': ext_out}, 
            files={'file': f}
        )

    if res_upload.status_code == 200:
        print('6. SUCCESS! Process complete.')
    else:
        raise Exception(f"Webhook upload failed. Status code: {res_upload.status_code}")

except Exception as e:
    err_str = str(e)
    print(f"CRITICAL ERROR during separation: {err_str}")
    report_failure(err_str)
    sys.exit(1)