import streamlit as st
from freeGPT import Client
import requests
import io
from PIL import Image
import os

api_token = os.environ.get("api_token")
API_URL = "https://api-inference.huggingface.co/models/goofyai/3d_render_style_xl"
headers = {"Authorization": f"Bearer {api_token}"}

st.title("Freak AI - Your stupid assistant")
st.header("Type any message and get stupid response", divider="rainbow")
messages = st.container(height=400)
prompt = st.chat_input(placeholder="Eg. How are you today?")

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.content

if prompt:
    if "Gerenate" in prompt or "Imagine" in prompt or "gemerate" in prompt or "imagine" in prompt or "Create" in prompt or "create" in prompt or "Image" in prompt:
        image_bytes = query({
    	    "inputs": prompt,
        })
        image = Image.open(io.BytesIO(image_bytes))

        user_msg = messages.chat_message("user", avatar="https://cdn.discordapp.com/attachments/1159192117751578705/1207020816332619796/kvikontent_professional_shot_of_clever_man_witha_book_in_hands__73d404c7-c515-4d49-b989-ac84dfe1e4ea.png?ex=65de20bb&is=65cbabbb&hm=1086ecaa2cfd107aa2912981226548f17bef10d54459f69d22176d38bdb1a05a&")
        user_msg.write(prompt)
        resp_msg = messages.chat_message("assistant", avatar="https://cdn.discordapp.com/attachments/1159192117751578705/1207019218466373642/kvikontent_professioanl_shot_of_stupid_man_c81e05bf-4530-497e-a9f3-63acdeea5345.png?ex=65de1f3f&is=65cbaa3f&hm=79516030ac8aa4d53f19c82a05cacd32ec0dceeb9b75d329a253bf4f3b8727a9&")
        resp_msg.image(image, caption=prompt)
    else:
        resp = Client.create_completion("gpt3", prompt + " . Give answer as very stupid dude guy, as stupid and funny as it can be. I should laugh when i'll read answer")
        user_msg = messages.chat_message("user", avatar="https://cdn.discordapp.com/attachments/1159192117751578705/1207020816332619796/kvikontent_professional_shot_of_clever_man_witha_book_in_hands__73d404c7-c515-4d49-b989-ac84dfe1e4ea.png?ex=65de20bb&is=65cbabbb&hm=1086ecaa2cfd107aa2912981226548f17bef10d54459f69d22176d38bdb1a05a&")
        user_msg.write(prompt)
        resp_msg = messages.chat_message("assistant", avatar="https://cdn.discordapp.com/attachments/1159192117751578705/1207019218466373642/kvikontent_professioanl_shot_of_stupid_man_c81e05bf-4530-497e-a9f3-63acdeea5345.png?ex=65de1f3f&is=65cbaa3f&hm=79516030ac8aa4d53f19c82a05cacd32ec0dceeb9b75d329a253bf4f3b8727a9&")
        resp_msg.write(resp)