from snowflake.snowpark import Session
from snowflake.snowpark.context import get_active_session
from snowflake.snowpark.exceptions import SnowparkSessionException
from snowflake.cortex import complete, CompleteOptions
import streamlit as st

from utils import (
    config_options,
    get_chat_history,
    init_messages,
    show_chat_history,
    feedback,
    generate_prompt,
    generate_first_prompt,
    get_content_instructions,
    llm_review,
)

try:
    session = get_active_session()
except SnowparkSessionException:
    session = Session.builder.config("connection_name", "email_draft_poc").create()
    # alternative approach
    # conn = st.connection("email_draft_poc", type="snowflake")
    # session = conn.session()

if "complete_feedback" not in st.session_state:
    st.session_state.complete_feedback = False

if "first_llm_response" not in st.session_state:
    st.session_state.first_llm_response = ""


config_options()
init_messages()

content_purpose = get_content_instructions(
    st.session_state.email_objective, st.session_state.single_album_tour_name
)
channel_limits = ""
if st.session_state.channel == "SMS":
    channel_limits = "Keep the content 3 sentence long!"

if st.session_state.complete_feedback:
    st.write("Thank you for your feedback!")

suggested_first_prompt = generate_first_prompt(
    st.session_state.artist, content_purpose, "", st.session_state.use_historical_emails
)

user_input = st.text_area(
    label="Suggested prompt for Generative AI based on initial selections from left side menu.",
    key="user_input",
    value=suggested_first_prompt,
    height=200,
)


# Create an Enter button
if st.button("Generate/modify email draft", type="primary"):
    st.session_state.complete_feedback = False
    st.session_state.first_content_created = True

    # Display the entered text when the button is clicked
    if user_input:
        chat_history = get_chat_history()

        llm_instructions = generate_prompt(
            user_input,
            st.session_state.first_llm_response,
            st.session_state.artist,
            st.session_state.email_objective,
            st.session_state.single_album_tour_name,
            st.session_state.first_prompt_by_llm,
            session,
            st.session_state.use_historical_emails,
        )
        with st.expander("Original/modified instructions to AI", expanded=False):
            st.code(llm_instructions)
        st.info("Final response from AI:")

        with st.spinner("Contacting our digital marketing guru...", show_time=True):
            options: CompleteOptions = {
                "temperature": st.session_state.temperature * 0.01,
                "guardrails": True,
            }
            initial_response = complete(
                model=st.session_state.model_name,
                prompt=llm_instructions,
                session=session,
                options=options,
            )

            response = llm_review(
                initial_response,
                st.session_state.llm_review_llm,
                st.session_state.artist,
                initial_response,
                session,
                st.session_state.first_llm_response,
            )
            st.session_state.first_llm_response = response

        st.write(response)
        st.session_state.messages.append({"role": "assistant", "content": response})
        st.session_state.messages.append({"role": "user", "content": llm_instructions})

        feedback(session)
    else:
        st.write("Please enter some text.")

st.divider()
show_chat_history(st.session_state.messages)