import os
import json
from flask import Flask, jsonify, request
from transformers import pipeline

# Create a Flask app
app = Flask(__name__)

# Initialize the audio model at the start of the API
audio_model = None

def download_models():
    global audio_model
    print("Downloading audio model...")
    # Download and load the audio model
    audio_model = pipeline("audio-classification", model="MelodyMachine/Deepfake-audio-detection-V2")
    print("Audio model downloaded and ready to use.")

# Download the model when the server starts
download_models()

@app.route('/detect', methods=['POST'])
def detect_deepfake():
    folder_path = request.form.get('folder_path')

    if not folder_path or not os.path.isdir(folder_path):
        return jsonify({"error": "Invalid folder path"}), 400

    results = {}
    try:
        # Process audio files only
        for file_name in os.listdir(folder_path):
            if file_name.endswith('.wav') or file_name.endswith('.mp3'):
                file_path = os.path.join(folder_path, file_name)
                result = audio_model(file_path)
                results[file_name] = {item['label']: item['score'] for item in result}

        # Save results to a file
        with open('detection_results.json', 'w') as f:
            f.write(json.dumps(results))

        return jsonify({"message": "Detection completed", "results": results}), 200

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    # Run the Flask app
    app.run(host='0.0.0.0', port=5000)