import streamlit as st
import os
import io
import PIL
import freeGPT

hf_token = os.environ.get("API_TOKEN")
API_URL = "https://api-inference.huggingface.co/models/Lykon/dreamshaper-xl-v2-turbo"
headers = {"Authorization": f"Bearer {hf_token}"}

st.title("🌎 AI Browser")
st.write("Browser that gives results using ai")
prompt = st.input_text("Search something", placeholder="Enter search query")
search_btn = st.button("Search")
results = st.container("Results")
texts, images = st.tabs(["All", "Images"])

def query(payload):
      response = requests.post(API_URL, headers=headers, json=payload)
      return response.content
        
def generate_image(prompt):
    image_bytes = query({
        "inputs": prompt,
    })
    image = Image.open(io.BytesIO(image_bytes))
    return image

if search_btn:
    prompt1 = prompt + ". Give short answer to this question"
    prompt2 = prompt + ". Give short funny answer to this question"
    prompt3 = prompt + ". Give short documental answer to this question"

    result1 = Client.create_completion("gpt3", prompt1)
    result2 = Client.create_completion("gpt3", prompt2)
    result3 = Client.create_completion("gpt3", prompt3)
    result4 = Client.create_completion("gpt3", prompt4)
    result5 = Client.create_completion("gpt3", prompt5)
    result6 = Client.create_completion("gpt3", prompt6)

    imageresult1 = generate_image(prompt + ". Create image of this in realistic style")
    imageresult2 = generate_image(prompt + ". Create image of this in anime style")
    imageresult3 = generate_image(prompt + ". Create image of this in black and white style")
    imageresult4 = generate_image(prompt + ". Create image of this in photo style")
    imageresult5 = generate_image(prompt + ". Create image of this in style of iphone low quality photo")
    imageresult6 = generate_image(prompt + ". Create image of this in cinemtic style")
    imageresult7 = generate_image(prompt + ". Create image of this in unreal engine style")
    imageresult8 = generate_image(prompt + ". Create image of this in google style")
    imageresult9 = generate_image(prompt + ". Create image of this in russia style")
    imageresult10 = generate_image(prompt + ". Create image of this in chinese style")
    imageresult11 = generate_image(prompt + ". Create image of this in style of sumsung low quality photo")
    imageresult12 = generate_image(prompt + ". Create image of this in romantic style")

    with texts:
        st.write("We found these results on your query: ")
        st.header(result1.split()[0], divider='rainbow')
        st.text(result1)
        st.header(result2.split()[0], divider='rainbow')
        st.text(result2)
        st.header(result3.split()[0], divider='rainbow')
        st.text(result3)
        st.header(result4.split()[0], divider='rainbow')
        st.text(result4)
        st.header(result5.split()[0], divider='rainbow')
        st.text(result5)
        st.header(result6.split()[0], divider='rainbow')
        st.text(result6)
        st.caption("That's the end!")

    with images:
        im1, im2, im3 = st.column(3)
        im4, im5, im6 = st.column(3)
        im9, im8, im7 = st.column(3)
        im10, im11, im12 = st.column(3)

        st.write("We found some images on your query: ")
        with im1:
            st.image(imageresult1)
        with im2:
            st.image(imageresult2)
        with im3:
            st.image(imageresult3)
        with im4:
            st.image(imageresult4)
        with im5:
            st.image(imageresult5)
        with im6:
            st.image(imageresult6)
        with im7:
            st.image(imageresult7)
        with im8:
            st.image(imageresult8)
        with im9:
            st.image(imageresult9)
        with im10:
            st.image(imageresult10)
        with im11:
            st.image(imageresult11)
        with im12:
            st.image(imageresult12)