import gradio as gr
from huggingface_hub import hf_hub_download
import importlib.util
import os
import sys

# Step 1: Dynamically load the model file
repo_id = "tlmk22/QuefrencyGuardian"
model_path = hf_hub_download(repo_id=repo_id, filename="model.py")
model_dir = os.path.dirname(model_path)
if model_dir not in sys.path:
    sys.path.append(model_dir)
spec = importlib.util.spec_from_file_location("model", model_path)
model_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(model_module)
FastModelHuggingFace = model_module.FastModelHuggingFace
fast_model = FastModelHuggingFace.from_pretrained(repo_id)

map_labels = {0: "chainsaw", 1: "environment"}  # Label mapping


def predict_audio(file):
    """
    Predict if a given audio file contains chainsaw activity or not.
    File: Input WAV file (uploaded via Gradio).
    """
    prediction = fast_model.predict(file, device="cpu")
    predicted_label = map_labels[prediction[0]]
    return f"Prediction: {predicted_label}"


example_files = [
    "/home/user/app/example1.wav",
    "/home/user/app/example2.wav",
]

# Build Gradio Interface

drag_and_drop_input = gr.Audio(type="filepath", label="Upload WAV File")
output_text = gr.Textbox(label="Prediction Result")

demo = gr.Interface(
    fn=predict_audio,
    inputs=drag_and_drop_input,
    outputs=output_text,
    examples=example_files,
    title="Quefrency Guardian: Chainsaw Noise Detector",
    description="Drag and drop a .wav audio file to predict whether it contains chainsaw noise or background environment sounds. Or record a 5-second audio clip and click the Record button to predict.",
)

if __name__ == "__main__":
    demo.launch()