import gradio as gr
import aiohttp
import asyncio
import json
import os

ENDPOINT_URL = os.getenv("ENDPOINT_URL")
WORKSPACE_TOKEN = os.getenv("WORKSPACE_TOKEN")
MAX_HISTORY_LENGTH = 10 

async def query_model(prompt, history):
    headers = {
        "Authorization": f"Bearer {WORKSPACE_TOKEN}",
        "Content-Type": "application/json"
    }
    
    limited_history = history[-MAX_HISTORY_LENGTH:]
    messages = []
    for msg in limited_history:
        messages.append({"role": "user", "content": msg[0]})
        if len(msg) > 1:
            messages.append({"role": "assistant", "content": msg[1]})
    
    messages.append({"role": "user", "content": prompt})
    
    data = json.dumps({
        "messages": messages
    })
    
    async with aiohttp.ClientSession() as session:
        async with session.post(ENDPOINT_URL, headers=headers, data=data) as response:
            if response.status == 200:
                response_json = await response.json()
                response_text = response_json[0]
            else:
                response_text = f"Error querying model: {response.status}"
    
    
    history.append((prompt, response_text))
    return history, history, ""


with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue=gr.themes.Color(
            c100="#30ea03",
            c200="#30ea03",
            c300="#30ea03",
            c400="#30ea03",
            c50="#30ea03",
            c500="#30ea03",
            c600="#30ea05",
            c700="#047857",
            c800="#30ea03",
            c900="#064e3b",
            c950="#054436"
        ),
        secondary_hue="slate",
        neutral_hue="zinc"
    ),
    css=""" 
    .gradio-container {
        background-color: #ecfdf5;  /* Light limegreen background color */
        padding: 20px;
        border-radius: 10px;
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
        display: flex;
        justify-content: center;
        align-items: center;
        height: 100vh;
        font-family: 'Protest Riot', sans-serif;  /* Custom font */
    }
    .gr-textbox {
        border: 2px solid #30ea03;
        border-radius: 8px;
        padding: 10px;
        color: black;  /* Set text color to black */
        font-family: 'Protest Riot', sans-serif;  /* Custom font */
    }
    .gr-button {
        background-color: #30ea03;
        color: white;
        border: none;
        border-radius: 8px;
        padding: 10px 20px;
        cursor: pointer;
        font-family: 'Protest Riot', sans-serif;  /* Custom font */
    }
    .gr-button:hover {
        background-color: #059669;
    }
    .gradio-app h1 {
        text-align: center;
    }
    .gradio-app .gr-interface-container .gr-interface-body .gr-interface-body-content .gr-interface-body-content-description {
        text-align: center;
    }
    .gradio-app .gr-interface-container .gr-interface-body .gr-interface-body-content .gr-inputs .gr-input .gr-label {
        color: white;  /* Change color of the input label to white */
    }
    .gradio-app .gr-interface-container .gr-interface-body .gr-interface-body-content .gr-outputs .gr-output .gr-label {
        color: white;  /* Change color of the output label to white */
    }
    .banner {
        text-align: center;
    }
    .banner img {
        width: 100px;  /* Adjust the width to improve image quality */
        height: auto;
    }
    #submit-button {
        background-color: #30ea03 !important;
        color: white !important;
        border: none !important;
        border-radius: 8px !important;
        padding: 10px 20px !important;
        cursor: pointer !important;
        font-family: 'Protest Riot', sans-serif !important;  /* Custom font */
    }
    #submit-button:hover {
        background-color: #059669 !important;
    }
    """
) as app:
    gr.Markdown("""
        <div class="banner">
            <img src="https://aadcdn.msauthimages.net/c1c6b6c8-iwpmhq8ytw8brx-8pyau9lwtqt1tlslp-rp7wl-zycs/logintenantbranding/0/bannerlogo?ts=637937709226757206" alt="Haleon Logo">
        </div>
    """)
    
    gr.Markdown("<h1 style='text-align: center;'>Haleon Expert Assist Tool</h1>")
    
    chatbot = gr.Chatbot(label="Chat History")
    state = gr.State([])  
    
    with gr.Row():
        with gr.Column():
            user_input = gr.Textbox(
                lines=2,
                placeholder="Enter your prompt here...",
                label="User Prompt",
                container=True,
                show_copy_button=True
            )
            submit_button = gr.Button("Submit", elem_id="submit-button")
    
    def submit_action(prompt, history):
        return asyncio.run(query_model(prompt, history))
    
    submit_button.click(fn=submit_action, inputs=[user_input, state], outputs=[chatbot, state, user_input])
    user_input.submit(fn=submit_action, inputs=[user_input, state], outputs=[chatbot, state, user_input])
    
app.queue(default_concurrency_limit=None) 
app.launch(share=False, show_error=True)