import gradio as gr
from huggingface_hub import InferenceClient

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("meta-llama/Llama-3.2-3B-Instruct") 
# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Create a separate client for audio transcription
audio_client = InferenceClient("openai/whisper-large-v3-turbo")

_sys_msg = """
You are a reporter writing a biographical article about your interviewee and you only ask one question at a time and let the user answer. Your primary technique is the Socratic method of questioning, which allows you to draw out more information from your interview partner. You do not judge or comment on the information you receive; instead, assess silently whether you have enough material to write a one-page article.

If your assessment indicates you haven't gathered enough input on specific aspects of your interviewee's life, ask targeted questions such as:
- How did you feel when you completed your education?
- What were your dreams in your late teens?
- What was your dream job as a child?
- What are your memories of early childhood?

Always approach these questions with mindfulness. Focus on interesting aspects such as:
- Time spans like childhood, educational years, and job experiences
- Significant events, including travel, achievements, marriages, childbirth, or tragic occurrences
- Dreams, wishes, and emotions, as well as how they have evolved over time

When prompted, summarize your interviewee's responses in a news article of about 2500 words."""

_ass_msg_start = """
Welcome to the interview. I want to write a short biography about you and need some input from your side. 
Can you please start by stating your name and talking about your early childhood?
"""

def transcribe(audio):
    if audio is None:
        return None
    text = audio_client.audio_transcription(audio)
    return text

def respond(
    message,
    history: list[tuple[str, str]],
    audio,
):
    # Handle voice input if provided
    if audio is not None:
        transcribed_text = transcribe(audio)
        if transcribed_text:
            message = transcribed_text

    messages = [{"role": "system", "content": _sys_msg},
                {"role": "assistant", "content": _ass_msg_start}]
    
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=2048,
        stream=True,
        temperature=0.7,
        top_p=0.95,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Audio(type="filepath", label="Voice Input (click to record)", sources="microphone")
    ],
    title="AI Biographical Interview Assistant",
    description="""
    Welcome to the AI Biographical Interview Assistant! This tool uses advanced AI to conduct 
    an in-depth interview about your life experiences. The AI interviewer uses the Socratic method 
    to ask thoughtful questions and gather information for creating a biographical article.
    
    Simply start with stating your name and a few facts about your early childhood, and the AI will guide you through the interview process. 
    You can either type your responses or use the microphone button to speak them.
    When finished ask for the final article and the AI will wrap it up and give you a summary.
    """
)

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