import os
from groq import Groq
import streamlit as st

client = Groq(
    api_key=os.environ.get("GROQ_API_KEY"),
)

st.write('Background analysis')
gl = st.text_area("Goal and assumptions",
                  placeholder="""What is the assumption to verify, what do you want to know,
                  the goal and mission of this polling""")
bg = st.text_area("Background information",
                  placeholder="""What is the polling about, what we should know to better
                  help you design the questions""")

if 'ft' in st.session_state or st.button('Next step'):
    if 'ft' not in st.session_state:
        rf = os.environ.get("DESIGN_PROMPT")
        messages = [
            {
                "role": "user",
                "content": f"""Generate possible factors related to the goal
                based on the information provided. List the factors in markdown list.
                You can give advice to human to edit their goal if you cannot define
                more than 3 factors.\n\n
                goals:\n\n{gl}\n\n
                background:\n\n{bg}\n\n
                some design principles for your reference:\n\n
                reference:\n\n```markdown\n{rf}\n\n```\n\n""",
            }
        ]
        
        chat_completion = client.chat.completions.create(
            messages=messages,
            model="llama3-70b-8192",
        )
        st.session_state['ft'] = chat_completion.choices[0].message.content
    
    am = st.text_area("Assumption to be tested",
                      placeholder="""What is the assumption you want to verify
                      through the polling""")
    st.session_state['ft'] = st.text_area(
        "Factors to be considered in questions",
        st.session_state['ft'],
        height=300)
    st.info('Before next step. Edit above factors, add or remove if needed')

    ft = st.session_state['ft']
    rf = os.environ.get("DESIGN_PROMPT")

    if st.button('Generate questions'):
        rf += f'''\n\n---\n\n{os.environ.get("DESIGN_PROMPT_SIXPOINT")}'''
        messages = [
            {
                "role": "user",
                "content": f"""Extract all factors from the text to create context.
                Generate 5 question based on the text and context.\n\n
                assumption to be tested through the questionnaire:\n\n{am}\n\n
                text:\n\n{ft}\n\n
                some design principles for your reference:\n\n
                reference:\n\n```markdown\n{rf}\n```\n\n""",
            }
        ]
        
        chat_completion = client.chat.completions.create(
            messages=messages,
            model="llama3-70b-8192",
        )

        md = chat_completion.choices[0].message.content
        st.markdown(md)
        st.info('Suggested number of sampling is 5 to 10 times of number of questions', icon="ℹ️")
        # dump questions into excel and allow download

        if st.button('Restart'):
            del st.session_state['ft']