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)

# Add downloaded path to sys.path for Python module recognition
if model_dir not in sys.path:
    sys.path.append(model_dir)

# Load the model dynamically
spec = importlib.util.spec_from_file_location("model", model_path)
model_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(model_module)

# Load FastModelHuggingFace class
FastModelHuggingFace = model_module.FastModelHuggingFace

# Step 2: Load the pre-trained model (dynamically from HuggingFace Hub)
fast_model = FastModelHuggingFace.from_pretrained(repo_id)

# Step 3: Define a prediction function
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")  # Assume CPU inference
    predicted_label = map_labels[prediction[0]]
    return f"Prediction: {predicted_label}"


# Step 4: Build Gradio Interface

# Define Gradio app elements
drag_and_drop_input = gr.Audio(type="filepath", label="Upload WAV File")
output_text = gr.Textbox(label="Prediction Result")

# Create Gradio Application
demo = gr.Interface(
    fn=predict_audio,
    inputs=drag_and_drop_input,
    outputs=output_text,
    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.",
)

# Launch App
if __name__ == "__main__":
    demo.launch()