import os
import json
from flask import Flask, jsonify, request
from transformers import pipeline
from pydub import AudioSegment
from io import BytesIO

# Create a Flask app
app = Flask(__name__)

# Initialize models at the start of the API
audio_model = None

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

# Download model when the server starts
download_models()

def preprocess_audio(file):
    # Load audio file
    audio = AudioSegment.from_file(file)
    
    # Convert audio to mono and normalize volume
    audio = audio.set_channels(1).set_frame_rate(16000)
    
    # Ensure audio is of a standard length (e.g., 10 seconds)
    duration_ms = len(audio)
    target_duration_ms = 10000  # Target duration in milliseconds (10 seconds)
    if duration_ms < target_duration_ms:
        # Pad with silence if shorter than target duration
        padding = AudioSegment.silent(duration=target_duration_ms - duration_ms)
        audio = audio + padding
    elif duration_ms > target_duration_ms:
        # Truncate if longer than target duration
        audio = audio[:target_duration_ms]

    return audio

@app.route('/detect', methods=['POST'])
def detect_deepfake():
    # Expect an audio file in the request
    audio_file = request.files.get('audio_file')

    # If a single audio file is provided
    if audio_file:
        try:
            # Preprocess the audio file
            audio = preprocess_audio(audio_file)
            
            # Save the processed file temporarily
            temp_wav = BytesIO()
            audio.export(temp_wav, format="wav")
            temp_wav.seek(0)

            # Perform detection
            result = audio_model(temp_wav)
            result_dict = {item['label']: item['score'] for item in result}

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

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

    # Invalid request if no audio file is provided
    else:
        return jsonify({"error": "Invalid input. Please provide an audio file."}), 400

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