import gradio as gr
import os
import subprocess
import tempfile
from pathlib import Path
from inference import AudioClassifier, Config

# Load classifier config
config = Config(
    yamnet_model_path="yamnet/yamnet.h5",
    yamnet_classes_path="yamnet/yamnet_class_map.csv",
    model_path="model/model.h5",
    custom_classes_path="model/model.npy",
    output_dir="results",
    output_file="classification.txt"
)
classifier = AudioClassifier(config)

NON_WAV_FORMATS = {".mp3", ".m4a", ".aac", ".flac", ".ogg", ".wma", ".aiff", ".ape", ".opus"}

def convert_to_wav(input_path: str) -> str:
    output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name
    cmd = ["ffmpeg", "-y", "-i", input_path, "-acodec", "pcm_s16le", "-ar", "44100", "-ac", "1", output_path]
    subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
    return output_path

def classify_audio(audio_path):
    messages = []

    ext = Path(audio_path).suffix.lower()
    if ext != ".wav":
        messages.append(f"🔁 Converting `{ext}` file to WAV format...")
        try:
            audio_path = convert_to_wav(audio_path)
            messages.append("✅ Conversion successful.")
        except Exception as e:
            return f"❌ Failed to convert file to WAV.\n\n**Error:** {str(e)}"

    messages.append("🔎 Running inference using YAMNet + fine-tuned model...")
    try:
        result = classifier.classify_file(audio_path)
    except Exception as e:
        return f"❌ Inference failed.\n\n**Error:** {str(e)}"

    label = result["dominant_label"]
    score = result["dominant_score_percentage"]
    messages.append(f"🎧 **Detected:** `{label}` with **{score}%** confidence.")

    return "\n\n".join(messages)

def upload_message(file):
    return (
        "📥 **Audio uploaded successfully!**\n\n"
        "▶️ You can now listen to the uploaded audio.\n"
        "🧠 Click **Submit** and Inference will run shortly."
    )

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🐝 AI-Belha Classifier\nClassify Queen Bee Status in beehives using YAMNet + Fine-Tuned Model.")
    
    with gr.Row():
        audio_input = gr.Audio(type="filepath", label="Upload audio file")
        output_box = gr.Markdown()

    audio_input.upload(fn=upload_message, inputs=audio_input, outputs=output_box)

    with gr.Row():
        submit_btn = gr.Button("🚀 Submit for Inference")
    
    submit_btn.click(fn=classify_audio, inputs=audio_input, outputs=output_box)

demo.launch()