import streamlit as st
import os

from init_env import init_env
from selenium_parser import extract_url, load_driver, scrape_page
from chatgpt import get_models, query


try:
    from dotenv import load_dotenv
    load_dotenv(".env")
    api_key = os.getenv("OPENAI_API_KEY")
except ImportError:
    api_key = None


@st.cache_data
def init():
    if api_key:
        print("LOCAL TEST")
    else:
        init_env()


@st.cache_resource
def get_driver():
    return load_driver()


def run():
    driver = get_driver()
    url_reviews = extract_url(url) if extract else url

    try:
        reviews = scrape_page(driver, url_reviews, page_count=page_count, wait_time=wait_time)
    except Exception as e:
        st.write(e)
        st.write("Page loaded:")
        st.image(driver.get_screenshot_as_png(), caption="screenshot")
        raise e
    
    st.write("Page loaded:")
    st.image(driver.get_screenshot_as_png(), caption="screenshot")
    with st.expander("Reviews"):
        st.json(reviews)
    st.markdown(f"Collected {len(reviews)} reviews")

    # ChatGPT answer
    stream = query(api_key, reviews, model=selected_model)

    st.divider()
    st.markdown(f"**{selected_model}:**")

    response = ""
    message_placeholder = st.empty()
    for chunk in stream:
        text = chunk.choices[0].delta.content
        if text is not None:
            response += text
            message_placeholder.markdown(response + "▌")
    message_placeholder.markdown(response)


init()

api_key = st.sidebar.text_input("Enter API key", value=api_key)
url = st.text_input("Enter URL", value="https://www.booking.com/hotel/th/queen-boutique.ru.html#tab-reviews")
extract = st.checkbox("Extract URL", value=True)

if api_key:
    models = get_models(api_key)
    selected_model = st.sidebar.selectbox("Select model", models)

st.sidebar.divider()
page_count = st.sidebar.slider("Number of pages", 0, 10, value=5)
wait_time = st.sidebar.number_input("Wait time (sec)", value=4)
    
if st.button("Run"):
    run()