from groq import Groq
import gradio as gr
from gtts import gTTS
import uuid
import base64
from io import BytesIO
import os

# Initialize Groq Client
# client = Groq(
#     api_key="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",
# )

client = Groq(api_key=os.getenv("GROQ_API_KEY_2"))

# client = Groq(
#     api_key="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",
# )

# Function to encode the image
def encode_image(uploaded_image):
    buffered = BytesIO()
    uploaded_image.save(buffered, format="PNG")  # Ensure the correct format
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

# Function to handle text and image inputs
def customLLMBot(user_input, uploaded_image):
    try:
        if uploaded_image is not None:
            # Encode the image to base64
            base64_image = encode_image(uploaded_image)

            # Create a message specifically for image prompts (no system messages allowed)
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "What's in this image?"},
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}},
                    ],
                }
            ]

            # Send the image message to the Groq API
            response = client.chat.completions.create(
                model="llama-3.2-11b-vision-preview",
                messages=messages,
            )
        else:
            # Process text input (system messages allowed here)
            messages = [
                {"role": "system", "content": "You are Dr. HealthBuddy, a professional virtual doctor chatbot."},
                {"role": "user", "content": user_input},
            ]
            response = client.chat.completions.create(
                model="llama-3.2-11b-vision-preview",
                messages=messages,
            )

        # Extract the reply
        LLM_reply = response.choices[0].message.content

        # Generate audio for response
        audio_file = f"response_{uuid.uuid4().hex}.mp3"
        tts = gTTS(LLM_reply, lang='en')
        tts.save(audio_file)

        return [(user_input or "Image uploaded", LLM_reply)], audio_file

    except Exception as e:
        # Handle errors gracefully
        error_message = f"An error occurred: {e}"
        return [(user_input or "Image uploaded", error_message)], None


# Gradio Interface
def chatbot_ui():
    with gr.Blocks() as demo:
        gr.Markdown("# Healthcare Chatbot Doctor")

        chatbot = gr.Chatbot(label="Responses")
        user_input = gr.Textbox(label="Ask a health-related question",
                                placeholder="Describe your symptoms...",
                                elem_id="user-input",
                                lines=1)
        uploaded_image = gr.Image(label="Upload an Image", type="pil")
        audio_output = gr.Audio(label="Audio Response")
        submit_btn = gr.Button("Submit")
        clear_btn = gr.Button("Clear")

        # Define actions
        submit_btn.click(
            customLLMBot,
            inputs=[user_input, uploaded_image],
            outputs=[chatbot, audio_output],
        )

        user_input.submit(
            customLLMBot,
            inputs=[user_input, uploaded_image],
            outputs=[chatbot, audio_output],
        )

        clear_btn.click(
            lambda: ([], "", None, []),
            inputs=[],
            outputs=[chatbot, user_input, uploaded_image, audio_output],
        )

    return demo

# Launch the interface
#chatbot_ui().launch(share=True)
chatbot_ui().launch(server_name="0.0.0.0", server_port=7860)


#chatbot_ui().launch(server_name="localhost", server_port=7860)