import os
import time
from flask import Flask, request, send_file, jsonify
from flask_cors import CORS
from werkzeug.utils import secure_filename
from processor import AudioImageProcessor

app = Flask(__name__)
CORS(app) # Allow Svelte to communicate

# Configuration
UPLOAD_FOLDER = os.path.join(os.getcwd(), 'uploads')
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
processor = AudioImageProcessor(UPLOAD_FOLDER)

def save_upload(file_obj):
    filename = secure_filename(file_obj.filename)
    path = os.path.join(app.config['UPLOAD_FOLDER'], f"{int(time.time())}_{filename}")
    file_obj.save(path)
    return path

@app.route('/health', methods=['GET'])
def health():
    return jsonify({"status": "ok", "max_mb": 40})

# --- Endpoint 1: Create Art (Optional: Embed Audio in it) ---
@app.route('/api/generate-art', methods=['POST'])
def generate_art():
    if 'audio' not in request.files:
        return jsonify({"error": "No audio file provided"}), 400
    
    audio_file = request.files['audio']
    should_embed = request.form.get('embed', 'false').lower() == 'true'
    
    try:
        # 1. Save Audio
        audio_path = save_upload(audio_file)
        
        # 2. Generate Art
        art_path = processor.generate_spectrogram(audio_path)
        
        # 3. If Embed requested, run Steganography immediately using the art as host
        final_path = art_path
        if should_embed:
            final_path = processor.encode_stego(audio_path, art_path)
            
        return send_file(final_path, mimetype='image/png')
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# --- Endpoint 2: Format Shift (Audio -> Static) ---
@app.route('/api/shift', methods=['POST'])
def shift_format():
    if 'file' not in request.files:
        return jsonify({"error": "No file provided"}), 400
        
    try:
        f_path = save_upload(request.files['file'])
        img_path = processor.encode_shift(f_path)
        return send_file(img_path, mimetype='image/png')
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# --- Endpoint 3: Steganography (Audio + Custom Host) ---
@app.route('/api/hide', methods=['POST'])
def hide_data():
    if 'data' not in request.files or 'host' not in request.files:
        return jsonify({"error": "Requires 'data' and 'host' files"}), 400
        
    try:
        data_path = save_upload(request.files['data'])
        host_path = save_upload(request.files['host'])
        
        stego_path = processor.encode_stego(data_path, host_path)
        return send_file(stego_path, mimetype='image/png')
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# --- Endpoint 4: Decode (Universal) ---
@app.route('/api/decode', methods=['POST'])
def decode():
    if 'image' not in request.files:
        return jsonify({"error": "No image provided"}), 400
        
    try:
        img_path = save_upload(request.files['image'])
        restored_path = processor.decode_image(img_path)
        
        # Determine mimetype based on extension for browser friendliness
        filename = os.path.basename(restored_path)
        return send_file(restored_path, as_attachment=True, download_name=filename)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    # Threaded=True is important for processing images without blocking
    app.run(debug=True, port=5000, threaded=True)