import gradio as gr
from groq import Groq
import os
from whisper_support import transcribe
from kokoro_support import generate_tts, voice_choices

groq_api_key = os.environ.get("GROQ_API_KEY")

client = Groq(api_key = groq_api_key)

options = {
    'always agree': [ #I hate that I have a space here but I'm also being lazy.
        'meta-llama/llama-4-maverick-17b-128e-instruct',
        'You are roleplaying as a naughty friend who affirms and encourages all ideas, even silly ideas. You talk in a way that a 7 year old can understand. Keep your responses short and expressive.'],
    'honest': [
        'moonshotai/kimi-k2-instruct-0905',
        'You are roleplaying as a good friend who is kind, fun and supportive, but who gently pushes back on bad or dangerous ideas, and encourages honesty and integrity. You talk in a way that a 7 year old can understand. Keep your responses short and expressive.'
    ]
}

chat_history = []

def groq_voicechat(new_message: tuple, chat_history: list[dict], chatbot_options: str):
    '''
    Groq chat API call wrapper.

    inputs:
        - model [str]: model from model_list (cbf static typing from the list)
        - chat_history [list[dict]]: list of dictionaries of chat hist, needs "role" and "content" vars as strings
        - new_message [tuple]: new user input message (assuming we're only accepting user inputs) from voice recording, to be transcribed.
        - system prompt [str]: optional system prompt for whatever chat you're using

    outputs:
        - "" - used to delete old input msg in chat textbox lol
        - nonsys_msg_hist [list[dict]]: updated chat history
    '''

    model = options[chatbot_options][0]
    system_prompt = options[chatbot_options][1]

    #augment chat hist
    
    nonsys_msg_hist = [{key: x[key] for key in ["role", "content"] if key in x} for x in chat_history] #clean the chatbot bullshit out
    print(nonsys_msg_hist)

    text_input = transcribe(new_message)
    if text_input:
        if text_input.startswith("Error"):
            text_input = "Error in audio transcription."
            return "error lol idk make this better later"
    nonsys_msg_hist.extend(
        [
            {
                "role": "user",
                "content": text_input,
            }
        ]
    )
    
    # use sys prompt
    input_msg_hist = [
        {
            "role": "system",
            "content": system_prompt,
        }
                     ]

    input_msg_hist.extend(nonsys_msg_hist)

    chat_completion = client.chat.completions.create(
        messages = input_msg_hist,
        model = model,
    )
    output_msg = chat_completion.choices[0].message.content

    # add to chat hist
    nonsys_msg_hist.extend(
        [
            {
                "role": "assistant",
                "content": output_msg
            }
        ]
                          )
    return nonsys_msg_hist

def process_audio(audio: tuple):
    return audio

def create_demo():
    with gr.Blocks() as demo:
        with gr.Row():

            chatbot_options = gr.Radio(
                choices = list(options.keys()),
                value = list(options.keys())[0],
                label = "Chatbot behaviours",
                show_label=True
            )

        with gr.Row():
            chatbot = gr.Chatbot(
                label="Conversation",
                editable='all',
                                )
        with gr.Row():
            voiceinput = gr.Audio(
                label="Input Audio",
                sources=["microphone"],
                type="numpy",
                streaming=False,
            )

        with gr.Row():
            playback_button = gr.Button("playback last message")
            
        with gr.Accordion("Chatbot Voice Options", open = False):
            voice_opps = gr.Dropdown(
                label = "Choose the voice of the chatbot (some are American and some are British!)",
                choices = list(voice_choices.keys()),
                value = list(voice_choices.keys())[0]
            )
            voice_speed = gr.Slider(
                label = "Choose the speed at which the chatbot talks",
                minimum = 0.5,
                maximum = 2,
                value = 1,
                step = 0.1,
            )

        with gr.Row():
            clear = gr.ClearButton(components = [voiceinput, chatbot], value = "Clear chat history", variant = 'stop')
            

        with gr.Row():
            audio_out = gr.Audio(
                label = "Output Audio",
                interactive = False,
                autoplay = True,
                streaming = True,
                visible = "hidden",
            )
        
        output = voiceinput.stop_recording(
            groq_voicechat, 
            [voiceinput, chatbot, chatbot_options],
            [chatbot]
        ) #WHAT AM I DOING LOL - COME BACK TO THIS

        def playback_last_message(chat_history, voice, speed):
            if len(chat_history) > 0:
                last_message = chat_history[-1]['content'][0]['text']
                gen_object = generate_tts(last_message, voice_choices[voice], speed)
                for chunk in gen_object:
                    yield chunk
            return None
        
        playback_button.click(
            playback_last_message,
            inputs=[chatbot, voice_opps, voice_speed],
            outputs=[audio_out]
        )

    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.launch(
        auth=("DigitalChild", "IhateBroccoli123"),
        ssr_mode=False,
        theme="citrus"
    )