import gradio as gr
from gradio_client import Client, handle_file
import os

client = Client(f"ahmedJaafari/{os.getenv('API_ENDPOINT')}", hf_token=os.getenv("API_KEY"))

def classify_audio(audio):
    # Simulating the classification result based on audio
    result = client.predict(
        audio=handle_file(audio),
        api_name="/predict"
    )
    transcript = result[1]
    service = result[0]
    
    # Define the color map for services
    service_color_map = {
        "Payment and Billing": "blue",
        "Technical Support": "green",
        "Customer Service": "purple",
        "Retention Service": "orange",
        "Other": "red"
    }
    
    # Check if the service is valid, else return an error message
    if service in service_color_map:
        service_color = service_color_map[service]
        service_html = f"<h1 style='color:{service_color}; font-weight:bold; text-align:center;'>{service}</h1>"
    else:
        service_html = "<h1 style='color:red; font-weight:bold; text-align:center;'>Error: No matching service found</h1>"
    
    return service_html, transcript

iface = gr.Interface(
    fn=classify_audio,
    inputs=[
        gr.Audio(type="filepath", label="Record Audio"),
    ],
    outputs=[
        gr.HTML(label="Service"),  # Colored HTML output for service
        gr.Textbox(label="Transcript")  # Transcript of the audio
    ],
    title="Vochai - IVR Routing System",
    description=(
        "This system helps route calls to the appropriate service based on audio input. "
        "The services include: Payment and Billing, Technical Support, Account Information, and Other.\n\n"
        "**Examples for each category:**\n"
        "- **Retention Service**: \"عفاك بغيت نحيد الكونطرا ديالي\"\n"
        "- **Payment and Billing**: \"بغيت نأجل أداء الفاتورة ديالي، شنو نقدر ندير؟\"\n"
        "- **Technical Support**: \"الإنترنت ما خدامش عندي، كيفاش نقدر نصاوبو؟\"\n"
        "- **Customer Support**: \"بغيت نعرف شنو هي تفاصيل الحساب ديالي.\"\n"
        "- **Other**: \"مافهمتش المشكل، شنو نقدر ندير؟\"\n\n"
        "Developed by Vochai. For more information, visit [Voch.ai](https://voch.ai) or contact us at contact@voch.ai.\n\n"
        "N.B. This Tool is for demonstration purposes only, the results are not 100% accurate and the speed is slower than usual."
    ),
)

iface.launch()