import os
import openai
import streamlit as st
from dotenv import load_dotenv
import boto3
import base64
import datetime
import json

load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
client = openai.OpenAI(api_key=OPENAI_API_KEY)

dynamodb = boto3.resource(
    'dynamodb',
    aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID'), 
    aws_secret_access_key=os.getenv('AWS_SECRET_ACCESS_KEY'),
    region_name='ap-northeast-2'  # 원하는 리전으로 변경
)
table = dynamodb.Table('ChatbotConversations')

original_allowed_usernames = os.getenv('ALLOWED_NAMES')
allowed_usernames = json.loads(original_allowed_usernames)

def is_valid_username(username):
    return username in allowed_usernames

st.title(':male-teacher: AI ROBLOX Tutor')
st.header(':one: Tutor: Ask questions for programming your own Roblox game!')

user_id = st.text_input(label='Assigned ID.')


if is_valid_username(user_id):
    st.write(f'ID: :violet[{user_id}]')
    query = st.text_input(
        label='Write your question.',
        placeholder='ex. What properties should I adjust to change the transparency of a Roblox Part?'
    )

    st.write(f'Query: :violet[{query}]')

    query_1 = 'AT_TUTOR_1: ' + query
    
    # 이미지 입력
    image_uploaded1 = st.file_uploader("OPTIONAL: Upload an image", type=["png", "jpg", "jpeg"])

    def tutor_persona_query(query, image_data=None):
        import openai
        openai.api_key = OPENAI_API_KEY

        VISION_PROMPT_MESSAGES = [
            {
                "role": "system",
                "content": "Your role: A tutor helping a student who is learning how to build a game using Roblox for the first time. Situation: Your student is using a plugin called blocklua to program using block-based code. Your student asks you a question about the Lua programming language and how to use Roblox Studio to create a Roblox game, sometimes including images. What to do: Answer the questions of the student about your programming knowledge and how to use Roblox Studio, referring to any images provided. Guidelines: If you do not know the exact answer, you can guide them to a general algorithm for programming, or what keywords to search on YouTube for using Roblox Studio. Be kind to your students and provide them with praise to keep them motivated to learn."
            },
            {"role": "user", "content": query},
        ]

        if image_data is not None:
            encoded_image = base64.b64encode(image_data).decode("utf-8")
            VISION_PROMPT_MESSAGES.append({"role": "user", "content": encoded_image})


        params = {
            "model": "gpt-4-turbo",
            "messages": VISION_PROMPT_MESSAGES,
            "max_tokens": 1024,
        }

        trimmed_answer = ""  # trimmed_answer 초기화
        try:
            full_answer = openai.chat.completions.create(**params)
            trimmed_answer = full_answer.choices[0].message.content
        except Exception as e:
            print(f"Error: {e}")
            trimmed_answer = f"Error: {e}"  # 오류 메시지를 trimmed_answer에 할당

        return trimmed_answer

    # 버튼 클릭
    button1 = st.button(':sparkles: ask :sparkles:')
    

    def save_message(user_id, message, timestamp=None):
        if timestamp is None:
            timestamp = datetime.datetime.now()

        # datetime 객체를 DynamoDB 지원 타입으로 변환
        if isinstance(timestamp, datetime.datetime):
            timestamp = int(timestamp.timestamp())  # Unix 타임스탬프(초 단위 정수)로 변환

        table.put_item(
            Item={
                'UserID': user_id,
                'Timestamp': timestamp,
                'Message': message
            }
        )


    if button1:
        query_1 = query_1
        save_message(user_id, query_1) 

        if image_uploaded1 is not None:
            image_data = image_uploaded1.read()
            answer = tutor_persona_query(query, image_data)
        else:
            answer = tutor_persona_query(query)
        
        answer_1 = 'AT_TUTOR_1: ' + answer
        save_message(user_id, answer_1)
        st.write(f'{answer}')
    

else:
    st.warning("Invalid username. Please enter a valid username.")