from typing import List, Dict, Any
import json

import pandas as pd
import streamlit as st
from snowflake.cortex import CompleteOptions, complete
from snowflake.cortex._complete import ResponseFormat
from snowflake.snowpark import Session
from concurrent.futures import ThreadPoolExecutor, as_completed

from previous_messages import (
    miley_facebook_messages,
    tate_facebook_messages,
    addison_facebook_messages,
    bsf_emails,
    russell_dickerson_emails,
    dave_east_emails,
    skye_newman_emails_messages,
)


def init_messages() -> None:
    """
    Initialize the session state messages and related variables.
    Clears the conversation if the 'clear_conversation' flag is set or if 'messages' is not in session state.
    """
    if st.session_state.clear_conversation or "messages" not in st.session_state:
        st.session_state.messages = []
        st.session_state.chat_history = ""
        # st.session_state.user_feedback_txt = ""
        # st.session_state.feedback_rating = ""
        st.session_state.first_llm_response = ""


# @st.dialog("Start over...")
def star_over_popup() -> None:
    """
    Function to reset Streamlit session state.
    """
    # st.success("Please follow these steps to start over and clear history/cache:")
    # st.write("1. From left top corner choose clear cache.")
    # st.write("2. Then refresh current webpage.")
    st.session_state.clear()
    st.rerun()


@st.dialog("Example artist messages")
def show_sample_messages(artist: str) -> None:
    """
    Trigger Streamlit pop up window showing sample messages used for context.
    """
    st.success("AI can use sample messages from artist to provide similar style.")
    if artist == "Miley Cyrus":
        st.write("Miley Cyrus Facebook messages:")
        miley_msg = miley_facebook_messages
        for i, msg in enumerate(miley_msg, 1):
            st.write(msg)
        st.divider()
    elif artist == "Addison Rae":
        st.write("Addison Rae Facebook messages:")
        addison_msg = miley_facebook_messages
        for i, msg in enumerate(addison_msg, 1):
            st.write(msg)
        st.divider()
    elif artist == "Tate McRae":
        st.write("Tate McRae Facebook messages:")
        tate_msg = miley_facebook_messages
        for i, msg in enumerate(tate_msg, 1):
            st.write(msg)
        st.divider()
    elif artist == "Russell Dickerson":
        st.write("Content from Russell Dickerson emails:")
        russell_dickerson_msg = russell_dickerson_emails
        for i, msg in enumerate(russell_dickerson_msg, 1):
            st.write(msg)
        st.divider()
    elif artist == "Dave East":
        dave_east_msg = dave_east_emails
        st.write("Content from Dave East emails:")
        for i, msg in enumerate(dave_east_msg, 1):
            st.write(msg)
        st.divider()
    elif artist == "Black Soprano Family":
        bsf_msg = bsf_emails
        st.write("Content from Black Soprano Family emails:")
        for i, msg in enumerate(bsf_msg, 1):
            st.write(msg)
        st.divider()
    elif artist == "Skye Newman":
        sn_msg = skye_newman_emails_messages
        st.write("Content from Skye Newman emails/facebook messages:")
        for i, msg in enumerate(sn_msg, 1):
            st.write(msg)
        st.divider()
    else:
        st.write("")


def config_options() -> None:
    """
    Configure the sidebar options for the Streamlit app.
    Allows the user to select a model, choose to remember chat history, provide assisting SQL documents,
    debug the summary of the previous conversation, and start over the session.
    """

    st.sidebar.selectbox(
        "Select channel",
        ("Email", "SMS"),
        key="channel",
        disabled=True,
    )
    st.sidebar.selectbox(
        "Artist sending out the email",
        (
            "Miley Cyrus",
            "Addison Rae",
            "Tate McRae",
            "Russell Dickerson",
            "Dave East",
            "Black Soprano Family",
            "Skye Newman",
        ),
        key="artist",
    )
    st.sidebar.selectbox(
        "Objective of the message",
        ("promote new album", "promote new single release", "promote merch", "promote tour"),
        key="email_objective",
    )
    st.sidebar.text_input(
        "Single/Album/Tour name", key="single_album_tour_name", value="Happy Memories"
    )
    st.sidebar.checkbox(
        "Include message samples from artist?",
        key="use_historical_emails",
        value=False,
    )

    if st.sidebar.button("Show example messages"):
        show_sample_messages(st.session_state.artist)
    st.sidebar.divider()
    if st.sidebar.button("Start Over"):
        star_over_popup()
    st.sidebar.divider()
    st.sidebar.checkbox(
        "Have AI generate first prompt",
        key="first_prompt_by_llm",
        value=False,
        help="AI will modify initial instructions for creating email content. Applies only to first instruction.",
        disabled=True,
    )
    st.sidebar.checkbox(
        "Have another AI polish the email content",
        key="llm_review_llm",
        value=False,
        help="AI will try to make initial email content more authentic and original",
        disabled=True,
    )
    st.sidebar.slider(
        "How creative you want email draft to be?",
        0,
        100,
        0,
        step=10,
        key="temperature",
        help="Higher-> Creative",
        disabled=True,
    )
    # https://docs.snowflake.com/user-guide/snowflake-cortex/aisql?ss_ad_code=SPA#choosing-a-model
    # https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api#model-availability
    st.sidebar.selectbox(
        "Select LLM:",
        (
            "claude-3-5-sonnet",
            "llama3.1-8b",
            "llama3.1-70b",
            "llama3.1-405b",
            "mistral-7b",
            "mistral-large",
            "mistral-large2",
        ),
        key="model_name",
        index=2,
        help="Models are shown from big to small",
        disabled=True,
    )
    # st.sidebar.checkbox(
    #     "Do you want to include Chartmetric metadata?",
    #     key="user_chartmetric_metadata",
    #     value=True,
    #     disabled=True,
    # )
    st.sidebar.button(".", key="clear_conversation", disabled=True)
    st.sidebar.expander("Session State").write(st.session_state)


def get_artist_history(artist: str) -> str:
    """
    Function to retrieve artist messages to be used as context for LLM prompt.
    """
    if artist == "Miley Cyrus":
        msgs = miley_facebook_messages
    elif artist == "Addison Rae":
        msgs = addison_facebook_messages
    elif artist == "Tate McRae":
        msgs = tate_facebook_messages
    elif artist == "Russell Dickerson":
        msgs = russell_dickerson_emails
    elif artist == "Dave East":
        msgs = dave_east_emails
    elif artist == "Black Soprano Family":
        msgs = bsf_emails
    elif artist == "Skye Newman":
        msgs = skye_newman_emails_messages
    else:
        msgs = []

    styled_output = "Here are some examples of artist previous messages:\n"

    for i, msg in enumerate(msgs, 1):
        styled_output += f"Example {i}:\n{msg.strip()}\n\n"

    return styled_output


def get_chat_history() -> List[Dict[str, Any]]:
    """
    Function to show formatted messages
    """
    slide_window = 5  # how many last conversations to remember. This is the slide window.
    chat_history = []

    start_index = max(0, len(st.session_state.messages) - slide_window)
    for i in range(start_index, len(st.session_state.messages) - 1):
        if not st.session_state.messages[i]["role"] == "assistant2":
            chat_history.append(st.session_state.messages[i])

    return chat_history


def show_chat_history(messages: List[Dict[str, Any]]) -> None:
    """
    Function to output formatted messages
    """
    counter = 0
    with st.expander("See chat history"):
        for message in reversed(messages):
            if message["role"] == "user":
                with st.chat_message("user"):
                    st.write(message["content"])
            elif message["role"] == "assistant":
                with st.chat_message("assistant"):
                    st.text(message["content"])
            else:
                with st.chat_message("ai"):
                    st.write(message["content"])
            counter += 1
            if counter % 2 == 0:
                st.divider()


def save_feedback(ses: Session) -> None:
    """
    Function to save user feedback into Snowflake table.
    """
    st.session_state.complete_feedback = True
    try:
        user = st.user.user_name
    except AttributeError:
        user = "localuser"

    feedback_df = pd.DataFrame(
        {
            "user": [user],
            "feedback_rating": [st.session_state.feedback_rating],
            "feedback_txt": [st.session_state.feedback_txt],
            "message_history": [st.session_state.messages],
            "first_prompt_by_llm": [st.session_state.first_prompt_by_llm],
            "email_objective": [st.session_state.email_objective],
            "llm_review_llm": [st.session_state.llm_review_llm],
            "use_historical_emails": [st.session_state.use_historical_emails],
            "artist": [st.session_state.artist],
        }
    )

    feedback_sp = ses.create_dataframe(data=feedback_df)
    feedback_sp.write.mode("append").save_as_table("email_draft_feedback", table_type="transient")


def feedback(ses: Session) -> None:
    """
    Code block to display feedback form with related button.
    """
    with st.form(key="my_form"):
        col1, col2 = st.columns(2)
        with col1:
            emotions = {
                "😁 Super satisfied": "super_Satisfied",
                "😊 Satisfied": "satisfied",
                "😐 Neutral": "neutral",
                "😞 Unsatisfied": "unsatisfied",
                "😭 Super unsatisfied": "super_unsatisfied",
            }

            # Create selectbox with emojis
            st.selectbox(
                "How happy are you with final result?", list(emotions.keys()), key="feedback_rating"
            )
        with col2:
            st.text_area("your feedback here", height=68, key="feedback_txt")
        st.form_submit_button(label="Submit evaluation", on_click=save_feedback, args=(ses,))


def generate_prompt(
    input: str,
    first_reply: str,
    artist: str,
    cta: str,
    single_album: str,
    first_prompt_from_ai: bool,
    ses: Session,
    prev_messages: bool,
    llm: str = "llama3.1-70b",
) -> str:
    """
    Function to create the very first prompt for LLM.
    If first_reply is set, then follow-up instructions from user will be applied.
    If first_prompt_from_ai is set, then AI will generate prompt instead of user input.
    """
    if prev_messages:
        context = get_artist_history(artist)
    else:
        context = ""

    if first_reply == "":
        instructions = f"""
        You task is to write prompt that another LLM can use to generate
        email draft. Email will be sent out by music artist. That another LLM
        should have persona definition reflecting somebody working with music marketing industry and with different artists.
        Prompt should include:
        A) instructions to keep in mind everything about the style that artist uses to communicate 
            with fans and try to replicate it in writing
        B) additional instructions based on the following known information:
        ARTIST_NAME: {artist}
        CALL_TO_ACTION: {cta}
        C) sample messages to provide additional context of the artist style:
        HISTORICAL_ARTIST_MESSAGES: {context}
        
        Output only the created prompt nothing else, no mention of the task itself.
        """
        if first_prompt_from_ai:
            options: CompleteOptions = {
                "temperature": 0,
                "guardrails": True,
            }
            response: str = complete(
                model=llm,
                prompt=instructions,
                session=ses,
                options=options,
            )
            return (
                response
                + "\nOnly output the email content, exclude comments how/why this kind of email was created."
            )
        else:
            return input
    else:
        prompt_with_history = f"""
        You are a digital marketing specialist working in a team of {artist} creating email content.
        \nYou have worked with various artists from music industry, so you know what kind of emails drive fan engagement.
        \nYou have received input from business side to make modifications to the original email.
        
        BUSINESS INPUT: {input}
        
        ORIGINAL EMAIL: {first_reply}.
        """

        return (
            prompt_with_history
            + "\nOnly output the email content, exclude comments how/why this kind of email was created."
        )


def generate_first_prompt(
    artist: str, purpose: str, channel_limits: str = "", history: bool = False
) -> str:
    """
    Function to show what first prompt would look like based on the selections from the side menu.
    """
    if history:
        context = get_artist_history(artist)
    else:
        context = ""

    first_prompt = f"""
    Imagine you are professional from email marketing team tasked with writing email content for {artist}. 
    \nYou have worked with various artists from music industry, so you know what kind of emails drive fan engagement.
    \nEmail should {purpose}.
    \nYour goal is to ensure created email content is engaging and attractive for typical audience of {artist}.
    Please keep in mind everything about the style that {artist} uses to communicate with fans and try to replicate it in your writing.
    Only output the email content, exclude comments how/why this kind of email was created.
    \n
    {context}
    """
    return first_prompt.strip()


def get_content_instructions(objective: str, release_name: str) -> str:
    """
    Function to retrieve more detailed call to action description based on the side menu input.
    """
    if objective == "promote merch":
        content_purpose = f"invite fans to purchase a merch related to the launch of a new album (let's call it {release_name})"
    elif objective == "promote new album":
        content_purpose = f"invite fans to open their favorite music app and listen to the newly released album (let's call it {release_name})"
    elif objective == "promote tour":
        content_purpose = f"invite fans to buy tickets to new tour (let's call it {release_name})"
    else:
        content_purpose = (
            f"invite fans to listen to the latest released track (let's call it {release_name})"
        )

    return content_purpose


# TODO ignore review if first reply already has been applied
def llm_review(
    response_for_review: str,
    review: bool,
    artist: str,
    content: str,
    ses: Session,
    first_content: str = "",
    llm: str = "llama3.1-70b",
) -> str:
    """
    Function to will have another different LLM review response and rewrite the email content.
    """
    review_prompt = f"""You are an expert from digital marketing agency. You have worked with various artists from music industry. 
        Your task is to review email content that will be sent to fans of {artist}. 
        \n
        How hard do you think it is for fans to guess that the above email is artificial, not handcrafted?
        If it seems too artificial then try to adjust the below email according to the suspicious you have.
        Think step by step before solving this. What do you know, what do you need to find, and how would you go about solving it?
        \n
        ```
        EMAIL_CONTENT: {content}
        ```
        \nOnly output new updated email content, nothing else about your suspicious or reasons for updating this email content.
    """

    if review and first_content == "":
        options: CompleteOptions = {
            "temperature": 0,
            "guardrails": True,
        }
        reviewed_response: str = complete(
            model=llm,
            prompt=review_prompt,
            session=ses,
            options=options,
        )
        return reviewed_response
    else:
        return content


def llm_verdict(content: str, session: Session, llm: str = "llama3.1-70b") -> Any:
    """
    Function to call LLM to extract details about email content.
    This can be used for understanding LLM response quality.
    """
    verdict_prompt = f"""You are a big music fan and you listen to different kind of artists.
     You have received and email and you have multiple tasks based on that email:
        1. Extract the name of the artist this email content is coming from. 
        \n
        2. Extract the call to action this email content is meant for.
        \n
        3. In the scale of 1 to 10 evaluate the email content itself based on how 
        original and non-artificial this email seems. Assign higher score to an email that looks more authentic, 
        and you believe is coming from actual music artist.
        \n
        4. Finally, in few sentences explain why you gave that score.
        \n
        ```
        EMAIL_CONTENT: {content}
        ```
        """

    response_format: ResponseFormat = {
        "type": "json",
        "schema": {
            "type": "object",
            "properties": {
                "verdict_for_content": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "artist_name": {"type": "string"},
                            "call_to_action": {  # TODO: use string only and calculate similarity score instead
                                "type": "string",
                                "enum": [
                                    "buy merch",
                                    "listen to track",
                                    "listen to album",
                                    "go to concert",
                                    "meet the artist",
                                    "visit artist website",
                                ],
                            },
                            "email_score": {
                                "type": "number",
                                "description": "Should be a number between 1 and 10",
                            },
                            "score_reason": {"type": "string"},
                        },
                        "required": [
                            "artist_name",
                            "call_to_action",
                            "email_score",
                            "score_reason",
                        ],
                    },
                }
            },
        },
    }
    # prompt = [{"role": "user", "content": "List 3 people and their ages"}]
    options: CompleteOptions = {
        "temperature": 0,
        "guardrails": True,
        "response_format": response_format,
    }

    reviewed_response: str = complete(
        model=llm,
        prompt=verdict_prompt,
        session=session,
        options=options,
    )
    return json.loads(reviewed_response)


def trigger_segment_based_emails(
    prompts: List[str], model: str, ses: Session, options: CompleteOptions
) -> list[str]:
    """
    Function to call multiple LLM requests in parallel and return those together.
    """

    def run_complete(
        prompt: str, llm_model: str, llm_ses: Session, llm_options: CompleteOptions
    ) -> str:
        updated_content_response: str = complete(
            model=llm_model,
            prompt=prompt,
            session=llm_ses,
            options=llm_options,
        )
        return updated_content_response

    results: list[str] = [""] * len(prompts)
    with ThreadPoolExecutor(max_workers=3) as executor:
        futures = {
            executor.submit(run_complete, prompt, model, ses, options): idx
            for idx, prompt in enumerate(prompts)
        }

        for future in as_completed(futures):
            idx = futures[future]
            try:
                result = future.result()
                results[idx] = result
            except Exception as e:
                print(f"Error in prompt {idx}: {e}")
                results[idx] = ""

    return results


def trigger_simple_llm_response(
    llm_model: str, llm_ses: Session, options: CompleteOptions, prompt: str
) -> str:
    """
    Function to call simple LLM
    """
    response: str = complete(
        model=llm_model,
        prompt=prompt,
        session=llm_ses,
        options=options,
    )
    return response