import streamlit as st
from snowflake.cortex import CompleteOptions
from snowflake.snowpark import Session
from snowflake.snowpark.context import get_active_session
from snowflake.snowpark.exceptions import SnowparkSessionException
from utils import trigger_simple_llm_response

try:
    session2 = get_active_session()
except SnowparkSessionException:
    session2 = Session.builder.config("connection_name", "email_draft_poc").create()


subject_line_sys_prompt = """You are email content creator in the music industry. Please suggest email subject line for
 the email. When suggesting subject line, take into a account the objective of the email. Recommend at most 5 
 subject lines.
\n
```
OBJECTIVE: [OBJECTIVE]
```
\n
```
EMAIL_CONTENT: [EMAIL_CONTENT]
```"""

with st.container(border=True):
    st.selectbox(
        "Objective of the message",
        ("promote new album", "promote new single release", "promote merch", "promote tour"),
        key="subject_email_objective",
    )
    st.text_area(
        label="Email content.",
        key="subject_email_content",
        height=200,
    )


subject_line_updated_sys_prompt = subject_line_sys_prompt.replace(
    "[OBJECTIVE]", st.session_state.subject_email_objective
).replace("[EMAIL_CONTENT]", st.session_state.subject_email_content)


if st.button("Generat subject line for email", type="primary"):
    with st.spinner("Contacting our digital marketing guru...", show_time=True):
        options: CompleteOptions = {
            "temperature": 0.5,
            "guardrails": True,
        }
        model = "llama3.1-8b"
        results = trigger_simple_llm_response(
            model, session2, options, subject_line_updated_sys_prompt
        )

        with st.expander("instructions to AI", expanded=False):
            st.code(subject_line_updated_sys_prompt)
        st.divider()
        st.write(results)