import streamlit as st
import random
import openai
import pandas as pd
import os

st.set_page_config(layout="wide")
st.balloons()


os.environ["OPENAI_API_KEY"]= st.secrets["OPENAI_API_KEY"]
openai.api_key = os.environ["OPENAI_API_KEY"]

if "df" not in st.session_state:
    st.session_state.df = pd.read_csv('matt.csv')

if "type" not in st.session_state:
    st.session_state.type = ''

if "time" not in st.session_state:
    st.session_state.time = ''

if "check" not in st.session_state:
    st.session_state.check = False

if "option" not in st.session_state:
    st.session_state.option = ''

if "result" not in st.session_state:
    st.session_state.result = []

if 'num_textareas' not in st.session_state:
    # 用来计数textarea的数量
    st.session_state['num_textareas'] = 1

def add_text_area():
    if st.session_state['num_textareas'] < 5:
        st.session_state['num_textareas'] += 1

def get_completion_from_messages(messages,
                                 model="gpt-3.5-turbo-16k",
                                 temperature=1.5, max_tokens=3000):
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return response.choices[0].message["content"]

def sample_generation(query, option = st.session_state.option, time = st.session_state.time, df = st.session_state.df):
    samples = df[option]
    print(samples)
    system_message = f'''
        You are an excellent press release writer;
        you craft narratives from a first-person perspective.
        You can reference the content and style of the examples
        to rewrite user input text in 【 】
        in the same style. Elaborate when necessary.
        Examples: {samples}
        '''

    messages =  [
                    {'role':'system',
                    'content': system_message + "keep it equal to {} words.".format(int(time)*60)},
                    {'role':'user',
                    'content': f"【{query}】"},]

    res = get_completion_from_messages(messages)
    return res

st.title('script generator')
tab1, tab2 = st.tabs(["Generation", "Library"])

with tab1:
    col1, col2 = st.columns([2, 3])
    with col1:
        st.session_state.option = st.selectbox(
    '',
    ("Intro", "News description", "Connection", 'Sponser', 'End of the video'))
        for i in range(st.session_state['num_textareas']):
            st.text_area(f"points {i+1}", key=f"textarea{i+1}")
        st.button("➕", on_click=add_text_area)
        st.session_state.time = st.slider("time mins", 0, 20, 5)
        if st.button('Generate'):
            # 可以在这里处理所有textarea的数据
            textarea_value = ''
            for i in range(st.session_state['num_textareas']):
                # 使用st.session_state获取特定的textarea输入值
                textarea_value += st.session_state.get(f"textarea{i+1}")

            progress_text = "Operation in progress. Please wait."
            my_bar = st.progress(0, text=progress_text)

            for i in range(2):
                my_bar.progress((i+1) * 30, text=progress_text)
                res = sample_generation(textarea_value)
                st.session_state.result.append(res)
            st.session_state.check = True
            my_bar.empty()

    with col2:
        if st.session_state.check:
            for i in range(len(st.session_state.result)):
                st.text_area("Sample" + str(i+1), value=st.session_state.result[i], height=300, key="text_area_"+str(i))
        
        else: st.image("https://static.streamlit.io/examples/dog.jpg")

with tab2:
    tab3, tab4, tab5, tab6, tab7 = st.tabs(["Intro", "News description", "Connection", 'Sponser', 'End of the video'])
    with tab3:
        st.table(st.session_state.df['Intro'])
    with tab4:
        st.table(st.session_state.df['News description'])
    with tab5:
        st.table(st.session_state.df['Connection'])
    with tab6:
        st.table(st.session_state.df['Sponser'])
    with tab7:
        st.table(st.session_state.df['End of the video'])