import eventlet
eventlet.monkey_patch()

import os
import requests
from flask import Flask, render_template, request, session, send_from_directory, jsonify
from flask_socketio import SocketIO, emit

app = Flask(__name__)
app.config['SECRET_KEY'] = 'docker_suno_final_v2'
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='eventlet')

MUSIC_FOLDER = os.path.join(os.getcwd(), 'saved_music')
if not os.path.exists(MUSIC_FOLDER):
    os.makedirs(MUSIC_FOLDER, exist_ok=True)

task_to_sid = {}

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/callback', methods=['POST'])
def callback():
    data = request.json
    if not data: return jsonify({"status": "no_data"}), 400
    
    task_id = data.get('taskId') or data.get('task_id')
    audio_url = data.get('audio_url')

    if audio_url and task_id:
        filepath = os.path.join(MUSIC_FOLDER, f"{task_id}.mp3")
        try:
            r = requests.get(audio_url, timeout=30)
            with open(filepath, 'wb') as f:
                f.write(r.content)
            
            if task_id in task_to_sid:
                socketio.emit('status', {
                    'msg': '✅ Saved to Server!',
                    'file_name': f"{task_id}.mp3",
                    'done': True
                }, room=task_to_sid[task_id])
        except Exception as e:
            print(f"Download Error: {e}")

    return jsonify({"status": "ok"}), 200

@app.route('/get_music/<filename>')
def get_music(filename):
    return send_from_directory(MUSIC_FOLDER, filename, as_attachment=True)

@socketio.on('save_token')
def save_token(data):
    session['api_token'] = data.get('token')
    emit('status', {'msg': '🔑 Token Linked.'})

@socketio.on('start_gen')
def handle_gen(payload):
    token = session.get('api_token')
    if not token:
        emit('status', {'msg': '❌ Error: Link token first.'})
        return

    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    
    # DOCKER FIX: Use the 'origin' header if available to set callback
    # If on HF, request.host_url works, but we provide a fallback
    callback_base = request.host_url.replace("0.0.0.0", "direct-url-here") 
    payload['callBackUrl'] = "https://zenjoul80-suno-api.hf.space/" + "callback"
    
    try:
        res = requests.post("https://api.sunoapi.org/api/v1/generate/upload-cover", json=payload, headers=headers)
        res_json = res.json()
        api_data = res_json.get('data')
        if api_data:
            task_id = api_data.get('taskId')
            task_to_sid[task_id] = request.sid
            emit('status', {'msg': f'🚀 Task {task_id} started...'})
        else:
            emit('status', {'msg': f"❌ API Error: {res_json.get('msg')}"})
    except Exception as e:
        emit('status', {'msg': f'❌ Connection Error: {str(e)}'})

if __name__ == '__main__':
    socketio.run(app, host="0.0.0.0", port=7860)