import gradio as gr
from huggingface_hub import InferenceClient
import whisper
from pydub import AudioSegment

# Load models
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
whisper_model = whisper.load_model("base")

# ---------- Whisper Audio-to-Text Function ----------
def transcribe_audio(file_path):
    try:
        print(f"📂 Processing audio: {file_path}")
        audio = AudioSegment.from_file(file_path)
        converted_path = "converted.wav"
        audio.export(converted_path, format="wav")
        result = whisper_model.transcribe(converted_path, fp16=False)
        return result["text"]
    except Exception as e:
        return f"❌ ERROR: {str(e)}"

# ---------- Step 1 ----------
def set_type(choice, user_profile):
    user_profile["interview_type"] = choice
    return "Great! What’s your background and what field/role are you aiming for?", user_profile

# ---------- Step 2 ----------
def save_background(info, user_profile):
    user_profile["field"] = info
    return "Awesome! Type 'start' below to begin your interview.", user_profile

# ---------- Step 3 ----------
def respond(message, chat_history, user_profile):
    if not user_profile.get("interview_type") or not user_profile.get("field"):
        bot_msg = "Please finish steps 1 and 2 before starting the interview."
        chat_history.append((message, bot_msg))
        return chat_history

    messages = [
        {"role": "system", "content": f"You are a professional interviewer conducting a {user_profile['interview_type']} interview for a candidate in the {user_profile['field']} field."}
    ]
    for q, a in chat_history:
        messages.append({"role": "user", "content": q})
        messages.append({"role": "assistant", "content": a})
    messages.append({"role": "user", "content": message})

    response = client.chat_completion(messages, max_tokens=150, stream=False)
    bot_msg = response.choices[0].message.content
    chat_history.append((message, bot_msg))
    return chat_history

# ---------- Handle Audio Input ----------
def handle_audio(audio_file, chat_history, user_profile):
    transcribed = transcribe_audio(audio_file)
    if transcribed.startswith("❌"):
        chat_history.append(("Audio input", transcribed))
        return chat_history
    return respond(transcribed, chat_history, user_profile)

# ---------- Gradio UI ----------
with gr.Blocks() as demo:
    user_profile = gr.State({"interview_type": "", "field": ""})
    chat_history = gr.State([])

    gr.Markdown("# 🎤 Welcome to Intervu")

    # Step 1: Choose Interview Type
    gr.Markdown("### Step 1: Choose Interview Type")
    with gr.Row():
        with gr.Column():
            btn1 = gr.Button("Behavioral")
            btn2 = gr.Button("Technical")
            btn3 = gr.Button("College / Scholarship")
        type_output = gr.Textbox(label="Bot response", interactive=False)

    btn1.click(set_type, inputs=[gr.Textbox(value="Behavioral", visible=False), user_profile], outputs=[type_output, user_profile])
    btn2.click(set_type, inputs=[gr.Textbox(value="Technical", visible=False), user_profile], outputs=[type_output, user_profile])
    btn3.click(set_type, inputs=[gr.Textbox(value="College", visible=False), user_profile], outputs=[type_output, user_profile])

    # Step 2: Enter Background
    gr.Markdown("### Step 2: Enter Your Background")
    background = gr.Textbox(label="Your background and field/goal")
    background_btn = gr.Button("Submit")
    background_output = gr.Textbox(label="Bot response", interactive=False)

    background_btn.click(save_background, inputs=[background, user_profile], outputs=[background_output, user_profile])

    # Step 3: Interview Chat
    gr.Markdown("### Step 3: Start Interview")
    chatbot = gr.Chatbot(label="Interview Bot")

    with gr.Row():
        msg = gr.Textbox(label="Your message")
        audio_input = gr.Audio(type="filepath", label="🎙️ Upload or Record your answer")

    with gr.Row():
        send_btn = gr.Button("Send Text")
        audio_btn = gr.Button("Send Audio")

    send_btn.click(respond, inputs=[msg, chat_history, user_profile], outputs=[chatbot], queue=False)
    send_btn.click(lambda: "", None, msg, queue=False)

    audio_btn.click(handle_audio, inputs=[audio_input, chat_history, user_profile], outputs=[chatbot], queue=False)

# ---------- Launch App ----------
demo.launch()