import requests
import gradio as gr
import json

def get_text_response(prompt):
    api_url = "http://35.233.231.20:5003/api/generate"
    data_payload = {
        "model": "llama3",
        "prompt": prompt,
        "stream": True
    }
    full_response = ""
    with requests.post(api_url, json=data_payload, stream=True) as response:
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8') 
                json_line = decoded_line.strip(',')
                try:
                    response_json = json.loads(json_line)
                    text_response = response_json.get("response", "")
                    done = response_json.get("done", False)
                    full_response += text_response 
                    yield full_response
                    if done:
                        break
                except json.JSONDecodeError as e:
                    continue 

def clear_chat():
    return "", ""

css = """
<style>
    .header {
        font-size: 24px;
        font-weight: bold;
        text-align: center;
        margin-top: 1rem;
        margin-bottom: 2rem;
    }
    button.gr-button {
        background-color: #7D3C98;
        color: white;
        padding: 0.5rem 1rem;
        font-size: 1rem;
        border-radius: 0.5rem;
    }
    button.gr-button:hover {
        background-color: #6C3483;
    }
    .acknowledgments {
        margin-top: 20px;
        text-align: center;
    }
    .gr-textbox, button.gr-button, .gr-output {
        height: 2.5rem;
    }
</style>
"""

with gr.Blocks(css=css) as demo:
    gr.HTML(
        """
            <div style="text-align: center; max-width: 650px; margin: 0 auto; padding-top: 7px;">
              <div
                style="
                  display: inline-flex;
                  align-items: center;
                  gap: 0.8rem;
                  font-size: 1.75rem;
                "
              >
                <h1 style="font-weight: 900; margin-bottom: 7px;">
                  Fair Compute Llama-3 Chat
                </h1>
              </div>
            </div>
        """
    )
    
    with gr.Row():
        prompt = gr.Textbox(label="Enter your prompt", lines=5, placeholder="Type something...")
        submit_button = gr.Button("Submit")
        clear_button = gr.Button("Clear")
    output = gr.Textbox(label="Response", lines=10, placeholder="Response will appear here...")

    submit_button.click(fn=get_text_response, inputs=prompt, outputs=output)
    clear_button.click(fn=clear_chat, inputs=[], outputs=[prompt, output])

    gr.HTML(
        """
        <div class="acknowledgments">
            <p>Run AI models on your home computers, powered by <a href="https://faircompute.com/" style="text-decoration: underline;" target="_blank">FairCompute</a></p>
        </div>
        """
    )

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