import streamlit as st
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain import FewShotPromptTemplate
from langchain.prompts.example_selector import LengthBasedExampleSelector
from dotenv import load_dotenv

load_dotenv() # load the env-sample.txt file

def getLLMResponse(query, age_option,tasktype_option):
    examples = []
    llm = OpenAI(temperature=.9, model="gpt-3.5-turbo-instruct")

#    if age_option=="Kid": #Silly and Sweet Kid 
#
#        examples = [
#        {
#            "query": "What is a mobile?",
#            "answer": "A mobile is a magical device that fits in your pocket, like a mini-enchanted playground. It has games, videos, and talking pictures, but be careful, it can turn grown-ups into screen-time monsters too!"
#        }
#        ]
#    
#    elif age_option=="Adult":  #Curious and Intelligent adult 
#        examples = [
#        {
#            "query": "What is a mobile?",
#            "answer": "A mobile is a portable communication device, commonly known as a mobile phone or cell phone. It allows users to make calls, send messages, access the internet, and use various applications. Additionally, 'mobile' can also refer to a type of kinetic sculpture that hangs and moves in the air, often found in art installations or as decorative pieces."
#        }
#        ]
#
#    elif age_option=="Senior Citizen": #A 90 years old guys
#        examples = [
#        {
#            "query": "What is a mobile?",
#            "answer": "A mobile, also known as a cellphone or smartphone, is a portable device that allows you to make calls, send messages, take pictures, browse the internet, and do many other things. In the last 50 years, I have seen mobiles become smaller, more powerful, and capable of amazing things like video calls and accessing information instantly."
#        }
#        ]


    example_template = """
    Question: {query}
    Response: {answer}
    """

    example_prompt = PromptTemplate(
        input_variables=["query", "answer"],
        template=example_template
    )


    prefix = """You are a {template_ageoption}, and you are going to {template_tasktype_option} , 
    you give one answer for each query. it is strictly limited to 1 answer only, and the answer MUST be LESS THAN 200 words. 
    For a tweet, you SHOULD NOT give more than 150 words. If it is not to write for a tweet, DO NOT give a tweet suggestion in your answer.
    """

    suffix = """
    Question: {template_userInput}
    Response: """

    example_selector = LengthBasedExampleSelector(
        examples=examples,
        example_prompt=example_prompt,
        max_length = numberOfWords
    )


    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(new_prompt_template.format(template_userInput=query,template_ageoption=age_option,template_tasktype_option=tasktype_option))
    response=llm(new_prompt_template.format(template_userInput=query,template_ageoption=age_option,template_tasktype_option=tasktype_option))
    print(response)

    return response

#UI Starts here

st.set_page_config(page_title="Marketing Tool",
                    page_icon='✅',
                    layout='centered',
                    initial_sidebar_state='collapsed')
st.header("Hey, How can I help you?")

form_input = st.text_area('Enter text', height=100)

tasktype_option = st.selectbox(
    'Please select the action to be performed',
    ('Write a sales copy', 'Create a tweet', 'Write a product description'),key=1)

age_option= st.selectbox(
    'For which target age group?',
    ('Kid' ,'Adult', 'senior Citizen'),key=2)

# numberOfWords= st.slider('Words limit', 1, 200, 25)
numberOfWords = 40 # the model doesn't take this.

submit = st.button("Generate")

if submit:
    st.write(getLLMResponse(form_input,tasktype_option,age_option))