import streamlit as st
from writer import write_article, incorporate_feedback, _template, evaluate_post
import hmac
import dspy
from dsp.modules import Claude

import phoenix as px

my_traces = px.Client().get_trace_dataset().save()
px.launch_app(trace=px.TraceDataset.load(my_traces))


def check_password():
    """Returns `True` if the user had the correct password."""

    def password_entered():
        """Checks whether a password entered by the user is correct."""
        if hmac.compare_digest(st.session_state["password"], st.secrets["password"]):
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # Don't store the password.
        else:
            st.session_state["password_correct"] = False

    # Return True if the password is validated.
    if st.session_state.get("password_correct", False):
        return True

    # Show input for password.
    st.text_input(
        "Password", type="password", on_change=password_entered, key="password"
    )
    if "password_correct" in st.session_state:
        st.error("😕 Password incorrect")
    return False


if not check_password():
    st.stop()


st.title("Linkedin shill")

#! I hate this
if "feedback_interface" not in st.session_state:
    st.session_state.feedback_interface = 0

    st.session_state.content = (
        ":heart_eyes: **Start writing today with the power of AI** :hugging_face:"
    )

with st.container(border=True):
    container = st.empty()

with st.container(border=True):
    container2 = st.empty()

container.markdown(st.session_state.content)

progress_text = "Operation in progress. Please wait."

with st.container():
    feedback_container = st.empty()

with st.sidebar:
    lm_choice = st.selectbox(
        "Select your language model",
        (
            "gpt4",
            "gpt-4-turbo",
            "gpt-3.5-turbo",
            "claude3 sonnet",
            "claude3 haiku",
            "claude3 opus",
            # "command-r-plus",
        ),
    )
    if lm_choice == "gpt-4-turbo":
        lm = dspy.OpenAI(
            model="gpt-4-0125-preview",
            max_tokens=3800,
            api_key=st.secrets["OpenAI"],
        )
    elif lm_choice == "gpt-3.5-turbo":
        lm = dspy.OpenAI(
            model="gpt-3.5-turbo",
            max_tokens=3800,
            api_key=st.secrets["OpenAI"],
        )
    elif lm_choice == "claude3 sonnet":
        lm = Claude(
            model="claude-3-sonnet-20240229",
            api_key=st.secrets["Claude"],
            max_tokens=4096,
        )
    elif lm_choice == "claude3 haiku":
        lm = Claude(
            model="claude-3-haiku-20240307",
            api_key=st.secrets["Claude"],
            max_tokens=4096,
        )
    elif lm_choice == "claude3 opus":
        lm = Claude(
            model="claude-3-opus-20240229",
            api_key=st.secrets["Claude"],
            max_tokens=4096,
        )
    # elif lm_choice == "command-r-plus":
    #     lm = LmChoice.command_r_plus
    else:
        lm = dspy.OpenAI(
            model="gpt-4",
            max_tokens=3800,
            api_key=st.secrets["OpenAI"],
        )

    with st.form("my_form"):
        topic = st.text_input("topic", "Oil future")
        template = st.text_area("template", _template, height=600)
        purpose = st.text_input("purpose", "Informative")
        audience = st.text_input("audience", "Linkedin")
        tone_style = st.text_input("tone style", "engaging and informative")
        key_points = st.text_input(
            "key points", "efficiency, climate change, energy dependance"
        )
        num_words = st.text_input("number of words", "600")
        language = st.text_input("language", "english")
        if st.form_submit_button(label="Submit"):
            user_inputs = {
                "topic": topic,
                "template": template,
                "purpose": purpose,
                "audience": audience,
                "tone_style": tone_style,
                "key_points": key_points,
                "num_words": num_words,
                "language": language,
            }
            bar = 1
            with st.status(progress_text) as status:
                container.status("Sit back and relax! AI is doing the job for you!")
                st.session_state.content = write_article(user_inputs, lm)
                st.session_state.feedback_interface = 1
                container.markdown(st.session_state.content)
                status.update(label="Writing complete!", state="complete")
                container2.status("AI is self evaluating!")
                evaluation = evaluate_post(st.session_state.content, lm)
                status.update(label="Evaluation complete!", state="complete")
                container2.markdown(evaluation)


if st.session_state.feedback_interface:
    with st.container(border=True):
        messages = st.container(height=300)
        if feedback := st.chat_input("What do you want to change?"):
            messages.chat_message("user").write(feedback)
            container.status("Sit back and relax! AI is doing the job for you!")
            st.session_state.content = incorporate_feedback(
                st.session_state.content,
                feedback,
                lm,
            )
            messages.chat_message("ai").write("Done!")
            container.markdown(st.session_state.content)