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

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

st.title("🌎 AI Browser")
st.write("Browser that gives results using ai. Warimg. It may took a while to generate answer")
prompt = st.text_input("Search something", placeholder="Enter search query")
search_btn = st.button("Search")
texts, images = st.tabs(["All results", "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 documentary answer to this question"
    prompt3 = prompt + ". Give short real life answer to this question"
    
    result1 = Client.create_completion("gpt3", prompt1)
    st.info("Result 1 is ready", icon='ℹ️')
    result2 = Client.create_completion("gpt3", prompt2)
    st.info("Result 2 is ready", icon='ℹ️')
    result3 = Client.create_completion("gpt3", prompt3)
    st.info("Result 3 is ready", icon='ℹ️')

    imageresult1 = generate_image(prompt + ". Create image of this in realistic style")
    st.info("Image result 1 is ready", icon='ℹ️')
    imageresult2 = generate_image(prompt + ". Create image of this in cinematic style")
    st.info("Image result 2 is ready", icon='ℹ️')
    imageresult3 = generate_image(prompt + ". Create image of this in black and white style")
    st.info("Image result 3 is ready", icon='ℹ️')
    st.info("starting showing output...", icon='ℹ️')

    with texts:
        st.write("We found these results on your query: ")
        
        st.header(result1.split()[0], divider='rainbow')
        st.caption("gptedia.com")
        st.text(result1)
        
        st.header(result2.split()[0], divider='rainbow')
        st.caption("reppit.com")
        st.text(result2)

        st.header(result3.split()[0], divider='rainbow')
        st.caption("llama-answers.com")
        st.text(result3)
        
        st.caption("That's the end!")

    with images:
        st.write("We found some images on your query: ")
        im1, im2, im3 = st.columns(3)
        with im1:
            st.image(imageresult1)
        with im2:
            st.image(imageresult2)
        with im3:
            st.image(imageresult3)