import gradio as gr
from groq import Groq
import base64
import asyncio
import os

client = Groq(
    api_key = os.getenv('GROQ_API_KEY')
)

STREAM_SPEED = 5

async def stream_text(text):
    words = text.split()
    result = ""
    delay = 0.1 / STREAM_SPEED
    
    for word in words:
        result += word + " "
        yield result
        await asyncio.sleep(delay)

def text_response(user_input):
    if not user_input.strip():
        return "Please enter some text to get a response."
        
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "system", 
                "content": "You are a helpfull assistant for a young aspiring Medical Student, Answer his queries correctly and always motivate him"
            },
            {"role": "user", "content": user_input}
        ],
        model="llama-3.2-11b-text-preview",
        temperature=0.5,
        max_tokens=2048,
        top_p=1,
        stop=None,
        stream=False,
    )

    return chat_completion.choices[0].message.content

def encode_image(image_path):
    with open(image_path, 'rb') as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def image_response(image_path, user_query):
    if not user_query.strip():
        return "Please enter some text along with the image."
    if image_path is None:
        return "Please upload an image to get a response."

    base64_image = encode_image(image_path)

    chat_completion = client.chat.completions.create(
        messages=[
            {"role": "user", "content": [
                {"type": "text", "text": user_query},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
            ]}
        ],
        model="llama-3.2-11b-vision-preview",
    )

    return chat_completion.choices[0].message.content

async def combined_response(text_input, image_input):
    if image_input is not None:
        response = image_response(image_input, text_input)
    else:
        response = text_response(text_input)
    
    async for streamed_response in stream_text(response):
        yield streamed_response

with gr.Blocks() as demo:
    with gr.Column():

        output_box = gr.Textbox(
            label="Assistant Response",
            placeholder="Response will appear here...",
            lines=5,
            interactive=False
        )
        
        text_input = gr.Textbox(
            label="Enter your message",
            placeholder="Type your message here...",
            lines=2,
        )
        
        image_input = gr.Image(
            label="Upload image (optional)",
            type="filepath"
        )
        
        submit_btn = gr.Button("Submit")
        
        text_input.submit(
            fn=combined_response,
            inputs=[text_input, image_input],
            outputs=output_box
        )

        submit_btn.click(
            fn=combined_response,
            inputs=[text_input, image_input],
            outputs=output_box
        )

demo.launch()