# -------------------------
# IMPORT REQUIRED LIBRARIES
# -------------------------

# Streamlit is a popular open-source framework used for building custom web apps for data science and ML.
import streamlit as st

# Custom libraries from langchain, for Few-Shot Learning Model interaction.
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain import FewShotPromptTemplate
from langchain.prompts.example_selector import LengthBasedExampleSelector

# Library for loading environment variables (for things like API keys).
from dotenv import load_dotenv

# -------------------------------
# LOAD ENVIRONMENT VARIABLES
# -------------------------------

# Load environment variables from a .env file.
load_dotenv()

# -------------------------------
# FUNCTION TO GET LLM RESPONSE
# -------------------------------

def getLLMResponse(query, age_option, tasktype_option):
    # Initialize the language model with specific settings.
    llm = OpenAI(temperature=.9, model="text-davinci-003")

    # We define different example sets based on the age group. These sets contain Q&A examples.
    
    # Examples for kids with fun and imaginative answers.
    if age_option == "Kid":
        examples = [ ... ] # List of child-friendly examples.

    # Thoughtful and elaborate answers tailored for adults.
    elif age_option == "Adult":
        examples = [ ... ] # List of adult-oriented examples.

    # Answers reflecting the wisdom and experiences of senior citizens.
    elif age_option == "Senior Citizen":
        examples = [ ... ] # List of examples for senior citizens.

    # Template for formatting the examples in the prompt.
    example_template = """
    Question: {query}
    Response: {answer}
    """

    # Define how our examples will be formatted using the PromptTemplate.
    example_prompt = PromptTemplate(
        input_variables=["query", "answer"],
        template=example_template
    )

    # The prefix sets up the model's persona and provides it with some example data.
    prefix = """You are a {template_ageoption}, and {template_tasktype_option}: 
    Here are some examples: 
    """
    
    # The suffix tells the model where to provide the answer.
    suffix = """
    Question: {template_userInput}
    Response: """

    # Example selector helps in selecting the best examples based on the given length.
    example_selector = LengthBasedExampleSelector(
        examples=examples,
        example_prompt=example_prompt,
        max_length=200
    )

    # This template combines everything: prefix, examples, and the suffix to create the full prompt.
    new_prompt_template = FewShotPromptTemplate(
        example_selector=example_selector,  # use example_selector instead of examples
        example_prompt=example_prompt,
        prefix=prefix,
        suffix=suffix,
        input_variables=["template_userInput", "template_ageoption", "template_tasktype_option"],
        example_separator="\n"
    )

    # Print the formatted prompt to the console (for debugging purposes).
    print(new_prompt_template.format(template_userInput=query, template_ageoption=age_option, template_tasktype_option=tasktype_option))
    
    # Fetch the response from the LLM using the prepared prompt.
    response = llm(new_prompt_template.format(template_userInput=query, template_ageoption=age_option, template_tasktype_option=tasktype_option))
    
    # Print the model's response to the console (for debugging purposes).
    print(response)

    # Return the response so it can be displayed on the Streamlit app.
    return response

# -------------------------
# STREAMLIT UI CONFIGURATION
# -------------------------

# Set the initial configurations for the Streamlit page (title, icon, layout).
st.set_page_config(page_title="Marketing Tool",
                   page_icon='✅',
                   layout='centered',
                   initial_sidebar_state='collapsed')

# Display a header on the web page.
st.header("Hey, How can I help you?")

# Create a text area where users can enter their query.
form_input = st.text_area('Enter text', height=275)

# Dropdown menu for selecting the type of task.
tasktype_option = st.selectbox(
    'Please select the action to be performed?',
    ('Write a sales copy', 'Create a tweet', 'Write a product description'), key=1)

# Dropdown menu for selecting the age group for the response.
age_option = st.selectbox(
    'For which age group is this intended?',
    ('Kid', 'Adult', 'Senior Citizen'), key=2)

# When the 'Submit' button is clicked, the entered query is processed.
if st.button('Submit'):

    # Call the `getLLMResponse` function to get the model's response.
    response = getLLMResponse(form_input, age_option, tasktype_option)

    # Display the response on the Streamlit page.
    st.write(response)

# ---------------------------------
# ADDITIONAL UI COMPONENTS
# ---------------------------------

# Display an information section on the page.
st.sidebar.info(
    "This tool is powered by the LangChain LLM and designed to provide tailored responses "
    "based on the selected age group and task type."
)

# Optional: Add any other UI components or information that might be useful for users.

# ------------------------------
# END OF STREAMLIT APPLICATION
# ------------------------------