import streamlit as st
from snowflake.cortex import CompleteOptions
from snowflake.snowpark import Session, FileOperation
from snowflake.snowpark.context import get_active_session
from snowflake.snowpark.exceptions import SnowparkSessionException
from snowflake.snowpark.functions import ai_complete, to_file
from utils import trigger_simple_llm_response
import os
from datetime import datetime

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


content_sys_prompt = """You are email content creator in the music industry. You know what emails
drive user engagement and increase email click and open ratios.
Please suggest 5 ways how to make this email more attractive. 
If given, take into consideration the feedback why this email hasn't been performing
well in the past or what it seems to be missing. Suggest at most 3 things to change.
\n
```
PREVIOUS_FEEDBACK: [PREVIOUS_FEEDBACK]
```
\n
```
EMAIL_CONTENT: [EMAIL_CONTENT]
```"""

content_as_image_sys_prompt = """You are email content creator in the music industry. You know what emails
drive user engagement and increase email click and open ratios.
Please suggest 5 ways how to make this email more attractive. 
If given, take into consideration the feedback why this email hasn't been performing
well in the past or what it seems to be missing. Suggest at most 3 things to change.
\n
```
PREVIOUS_FEEDBACK: [PREVIOUS_FEEDBACK]
```
\n
```
EMAIL_CONTENT: [EMAIL_CONTENT]
```"""


content_type = st.radio(
    "Specify email content type?",
    ["Email content as a text", "Email as a screenshot"],
)

with st.container(border=True):
    if content_type == "Email content as a text":
        st.text_area(
            label="Email content.",
            key="feedback_email_content",
            height=200,
        )
    else:
        st.write("upload screenshot:")
        image_file = st.file_uploader("Choose a file", type=["jpg", "jpeg", "png"])
        if image_file is not None:
            st.image(image_file, caption="Preview")
    st.text_area(
        label="Feedback to AI regarding this email content.",
        key="previous_email_feedback",
        height=100,
    )

if st.button("Review content", type="primary"):
    with st.spinner("Contacting our digital marketing guru...", show_time=True):
        options: CompleteOptions = {
            "temperature": 0.01,
            "guardrails": True,
        }
        if content_type == "Email content as a text":
            updated_content_sys_prompt = content_sys_prompt.replace(
                "[PREVIOUS_FEEDBACK]", st.session_state.previous_email_feedback
            ).replace("[EMAIL_CONTENT]", st.session_state.feedback_email_content)
            model = "claude-3-5-sonnet"
            model = "llama3.1-8b"
            results = trigger_simple_llm_response(
                model, session2, options, updated_content_sys_prompt
            )

            with st.expander("instructions to AI", expanded=False):
                st.code(updated_content_sys_prompt)
        else:
            updated_content_as_image_sys_prompt = content_as_image_sys_prompt.replace(
                "[EMAIL_CONTENT]", st.session_state.previous_email_feedback
            )

            IMAGES_STAGE = "EMAILPOC_IMAGES_STAGE"
            session2.sql(
                f"CREATE STAGE IF NOT EXISTS {IMAGES_STAGE} ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')"
            ).collect()

            FILE_PATH = (
                os.path.splitext(image_file.name)[0]
                + "-"
                + datetime.now().strftime("%Y-%m-%d_%H.%M.%S")
            )
            FINAL_FILE_PATH = "@" + IMAGES_STAGE + "/" + FILE_PATH + "/" + image_file.name
            FileOperation(session2).put_stream(
                input_stream=image_file,
                auto_compress=False,
                stage_location=FINAL_FILE_PATH,
            )

            df = session2.range(1).select(
                ai_complete(
                    model="claude-3-5-sonnet",
                    prompt=updated_content_as_image_sys_prompt,
                    file=to_file(FINAL_FILE_PATH),
                )
            )
            results = df.collect()[0][0]

            # results = ai_complete(
            #     model='claude-3-5-sonnet',
            #     prompt=updated_content_as_image_sys_prompt,
            #     file=FINAL_FILE_PATH,
            #     session=session2,
            #     options=options,
            # )
        st.divider()
        st.write(results)