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

# --- Variables injected by Flask ---
AUDIO_TYPE = '___AUDIO_TYPE___'
NEGATIVE_PROMPT = '''___NEGATIVE_PROMPT___'''

prompt = os.environ.get('PROMPT', '')
try: seed = int(os.environ.get('SEED', '-1'))
except: seed = -1
try: duration = float(os.environ.get('DURATION', '8'))
except: duration = 8.0

file_url = os.environ.get('FILE_URL', '')
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": "generate-audio",
                "client_payload": {
                    "data": {
                        "prompt": prompt,
                        "duration": duration,
                        "seed": seed,
                        "audio_type": AUDIO_TYPE,
                        "negative_prompt": NEGATIVE_PROMPT,
                        "file_url": file_url,
                        "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}")

try:
    print('1. Connecting to Audio Space...')
    client = Client("hkchengrex/MMAudio")
    
    if AUDIO_TYPE == 'video':
        print('2. Downloading video for audio generation...')
        req = requests.get(file_url, timeout=30)
        ext = file_url.split('.')[-1] if '.' in file_url else 'mp4'
        local_vid = f'input.{ext}'
        with open(local_vid, 'wb') as f:
            f.write(req.content)
            
        print('3. Processing Video-to-Audio...')
        result = client.predict(
            video=handle_file(local_vid),
            prompt=prompt,
            negative_prompt=NEGATIVE_PROMPT,
            seed=seed,
            num_steps=25,
            cfg_strength=4.5,
            duration=duration,
            api_name="/video_to_audio"
        )
    else:
        print('2. Processing Text-to-Audio...')
        result = client.predict(
            prompt=prompt,
            negative_prompt=NEGATIVE_PROMPT,
            seed=seed,
            num_steps=25,
            cfg_strength=4.5,
            duration=duration,
            api_name="/text_to_audio"
        )
        
    print('4. Preparing File Data...')
    audio_path = None
    if isinstance(result, list) and len(result) > 0:
        audio_path = result[0].get('video') or result[0].get('path') or result[0].get('url') if isinstance(result[0], dict) else result[0]
    elif isinstance(result, tuple) and len(result) > 0:
        audio_path = result[0]
    elif isinstance(result, dict):
        audio_path = result.get('video') or result.get('path') or result.get('url')
    else:
        audio_path = result
        
    ext = 'mp4'
    if str(audio_path).endswith('.wav'): ext = 'wav'
    if str(audio_path).endswith('.mp3'): ext = 'mp3'
    
    print('5. Uploading back to Space...')
    with open(audio_path, 'rb') as f:
        res = requests.post(
            f'{space_url}/api/webhook/upload', 
            data={'run_id': run_id, 'github_run_id': github_run_id, 'ext': ext}, 
            files={'file': f}
        )
        
    if res.status_code == 200:
        print('SUCCESS!')
    else:
        sys.exit(1)
        
except Exception as e:
    err_str = str(e)
    print(f'Generation error: {err_str}')
    report_failure(err_str)
    sys.exit(1)