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_segment_based_emails

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


newsletter_txt = """Calling all my current and former fans.
I am releasing my new album next week. Please hit this link
to presave and partcipate in signed T-shirt Giveway.
Yours truly!
"""
segment_based_user_input = st.text_area(
    label="Original newsletter/ email to be used for customization.",
    key="segment_based_user_input",
    value=newsletter_txt,
    height=100,
)

sys_prompt = """You are email content creator in the music industry. Please modify newsletter to meet
characteristics of target audience and make newsletter more appealing for them.
\n
```
NEWSLETTER: [NEWSLETTER]
```
\n
```
AUDIENCE: [SEGMENT_DESCRIPTION]
```"""

segment_based_system_prompt = st.text_area(
    label="Main instructions for AI.",
    key="segment_system_prompt",
    value=sys_prompt,
    height=100,
)

casual_fans, engaged_fans, super_fans = st.columns(3, vertical_alignment="bottom")

casual_characteristics = """This segment usually signs up for newsletters."""
with casual_fans:
    st.text_area(
        label="Main characteristics of casual fans.",
        key="casual_fans_description",
        value=casual_characteristics,
        height=100,
    )

engaged_characteristics = """This segment usually opens emails and clicks on links. They  also listen to music on Apple Music and watch YouTube videos."""
with engaged_fans:
    st.text_area(
        label="Main characteristics of engaged fans.",
        key="engaged_fans_description",
        value=engaged_characteristics,
        height=100,
    )

super_characteristics = """This segment very actively listens to music on different music platforms.
They will take advantage of chance to buy merch or concert tickets."""
with super_fans:
    st.text_area(
        label="Main characteristics of super fans.",
        key="super_fans_description",
        value=super_characteristics,
        height=100,
    )

if st.button("Generate/modify email draft for segments", type="primary"):
    with st.spinner("Contacting our digital marketing guru...", show_time=True):
        casual_sys_prompt = sys_prompt.replace("[NEWSLETTER]", newsletter_txt).replace(
            "[SEGMENT_DESCRIPTION]", casual_characteristics
        )
        engaged_sys_prompt = sys_prompt.replace("[NEWSLETTER]", newsletter_txt).replace(
            "[SEGMENT_DESCRIPTION]", engaged_characteristics
        )
        super_sys_prompt = sys_prompt.replace("[NEWSLETTER]", newsletter_txt).replace(
            "[SEGMENT_DESCRIPTION]", super_characteristics
        )

        prompts = [casual_sys_prompt, engaged_sys_prompt, super_sys_prompt]
        options: CompleteOptions = {
            "temperature": 0.01,
            "guardrails": True,
        }
        model = "llama3.1-8b"
        results = trigger_segment_based_emails(prompts, model, session2, options)

        casual_tab, engaged_tab, super_tab = st.tabs(
            ["Casual fan email", "Engaged fan email", "Super fan email"]
        )
        with casual_tab:
            with st.expander("instructions to AI", expanded=False):
                st.code(casual_sys_prompt)
            st.divider()
            st.write(results[0])
        with engaged_tab:
            with st.expander("instructions to AI", expanded=False):
                st.code(engaged_sys_prompt)
            st.divider()
            st.write(results[1])
        with super_tab:
            with st.expander("instructions to AI", expanded=False):
                st.code(super_sys_prompt)
            st.divider()
            st.write(results[2])