import gradio as gr
from transformers import pipeline
from langdetect import detect
import base64
import io
from pydub import AudioSegment

# ----------------------------
# Load pretrained model globally (fixes runtime error)
# ----------------------------
model = pipeline("audio-classification", model="microsoft/wavlm-base-sv")  # AI vs Human

# ----------------------------
# Function to detect language
# ----------------------------
def detect_language(audio_file):
    try:
        # Convert audio to text using Whisper (optional)
        from transformers import pipeline as hf_pipeline
        asr = hf_pipeline("automatic-speech-recognition", model="openai/whisper-small")
        transcription = asr(audio_file)["text"]
        lang = detect(transcription)
        return lang
    except Exception:
        return "unknown"

# ----------------------------
# Main classification function
# ----------------------------
def classify_audio(input_audio):
    try:
        # If input is Base64 string, decode it
        if isinstance(input_audio, str):
            audio_bytes = base64.b64decode(input_audio)
            audio_file = io.BytesIO(audio_bytes)
        else:
            audio_file = input_audio  # If it's a file from Gradio audio recorder

        # Convert to WAV for model
        audio_segment = AudioSegment.from_file(audio_file)
        wav_io = io.BytesIO()
        audio_segment.export(wav_io, format="wav")
        wav_io.seek(0)

        # Run AI vs Human model
        result = model(wav_io)[0]
        classification = "AI_GENERATED" if "AI" in result['label'].upper() else "HUMAN"
        confidence = round(float(result['score']), 2)

        # Detect language
        wav_io.seek(0)
        language = detect_language(wav_io)

        # Return JSON for buildathon
        return {
            "classification": classification,
            "language": language,
            "confidence": confidence
        }
    except Exception as e:
        return {"error": str(e)}

# ----------------------------
# Gradio Interface
# ----------------------------
with gr.Blocks() as demo:
    gr.Markdown("## AI vs Human Voice Detector + Language Detection")
    with gr.Row():
        audio_input = gr.Audio(source="microphone", type="filepath", label="Speak or upload audio")
        output = gr.JSON(label="Prediction")
    classify_btn = gr.Button("Classify")
    classify_btn.click(classify_audio, inputs=audio_input, outputs=output)

demo.launch()