from flask import Flask, request, jsonify
import whisper
import os

app = Flask(__name__)

# Load the Whisper model
print("Loading Whisper model...")
model = whisper.load_model("tiny")
print("Whisper model loaded.")

def transcribe(audio_path):
    print(f"Transcribing audio from: {audio_path}")

    # Load audio and pad/trim it to fit 30 seconds
    print("Loading and processing audio...")
    audio = whisper.load_audio(audio_path)
    audio = whisper.pad_or_trim(audio)

    # Make log-Mel spectrogram and move to the same device as the model
    print("Creating log-Mel spectrogram...")
    mel = whisper.log_mel_spectrogram(audio).to(model.device)

    # Detect the spoken language
    print("Detecting language...")
    _, probs = model.detect_language(mel)
    language = max(probs, key=probs.get)
    print(f"Detected language: {language}")

    # Decode the audio
    print("Decoding audio...")
    options = whisper.DecodingOptions(fp16=False)
    result = whisper.decode(model, mel, options)
    
    print("Transcription complete.")
    return result.text, language

@app.route('/transcribe', methods=['POST'])
def transcribe_audio():
    print("Received request at /transcribe")
    if 'audio' not in request.files:
        print("Error: No audio file provided")
        return jsonify({"error": "No audio file provided"}), 400
    
    audio_file = request.files['audio']
    
    # Save the uploaded audio file
    audio_path = os.path.join("temp_audio", audio_file.filename)
    os.makedirs("temp_audio", exist_ok=True)
    audio_file.save(audio_path)
    print(f"Audio file saved to: {audio_path}")
    
    # Transcribe the audio
    try:
        transcription, language = transcribe(audio_path)
    except Exception as e:
        print(f"Error during transcription: {str(e)}")
        return jsonify({"error": f"An error occurred: {str(e)}"}), 500
    
    # Clean up the saved file
    os.remove(audio_path)
    print(f"Audio file removed from: {audio_path}")
    
    # Return the transcription and detected language
    print(f"Transcription: {transcription}, Language: {language}")
    return jsonify({"transcription": transcription, "language": language}), 200

@app.route('/healthcheck', methods=['GET'])
def healthcheck():
    print("Received request at /healthcheck")
    return jsonify({"status": "API is running"}), 200

if __name__ == '__main__':
    print("Starting Flask app...")
    app.run(host="0.0.0.0", port=5000)