# ==============================================================================
# KU HMC Lab Chatbot
# Author: Kyle VanSaun, kavansaun@ku.edu
# Editors:
# Last Updated: April 2026
# Description: A ChatBot that can
#              mimic human behavior
#              and saves participant
#              responses to Github.
# Uses: Hugging Face as the Web host 
#       Gradio for the User Interface
#       Python for the code
#       OpenAI for the ChatBot
#       Github to save data
# ==============================================================================
# INSTRUCTIONS for Chat Bot on line 92, Insert Repository name on line 39, Model type on line 150

# Imports
import gradio as gr
from openai import OpenAI
import os
from github import Github
import uuid
import time

# Environment variables
openai_key = os.getenv("OPENAI_KEY")
github_key = os.getenv("GITHUB_KEY")

# Check if openai key was given
if not openai_key:
    raise gr.Error("OPENAI_KEY not set")
client = OpenAI(api_key=openai_key)

# Check if Github key was given
g = Github(github_key) if github_key else None

# Repo name, if you want to save to a github repository change RepositoryName to your repo name, name must be under quotations
REPO_NAME = "RepositoryName"

# Session data dictionary for each user
user_dictionary = {}

# Function to create or update the GitHub chat log
def update_github_log(uid):
    if g:
        # Try to get repo
        try:
            repo = g.get_user().get_repo(REPO_NAME)
            filename = f"{uid}.txt"
            content = ""
            for msg in user_dictionary[uid]:
                role = msg["role"].capitalize()
                content += f"{role}: {msg['content']}\n"
            
            # Try to update the repo, otherwise create the repo
            try:
                file = repo.get_contents(filename)
                repo.update_file(file.path, "Update chat log", content, file.sha)
            except Exception:
                repo.create_file(filename, "Create session log", content)
        # If we didnt get the repo find and show the error
        except Exception as e:
            error_msg = str(e).lower()
            if "bad_credentials" in error_msg or "bad credentials" in error_msg:
                gr.Warning("The GitHub API Key provided is invalid.")
            elif "repository" in error_msg or "404" in error_msg:
                gr.Warning("Repo not found, Possible incorrect Repository name")
            else:
                gr.Warning(f"GitHub Error: {str(e)}")
    else:
        pass

# Gradio user interface, status and progress made invisible
with gr.Blocks(theme=gr.themes.Monochrome(), css="""footer {display:none !important;}[class*="status"] {display:none !important;}[class*="progress"] {display:none !important;}""") as chatblock:
    
    # Hidden session variables
    user_id = gr.Textbox(visible=False)
    notifier = gr.HTML(visible=False) # If you embed this in html, this notifier is so you can catch and use the user_id in your script or url.
    
    # Load user session
    @chatblock.load(outputs=[user_id, notifier])
    def load_user():
        global user_dictionary
        
        # Creating user ID
        new_id = str(uuid.uuid4())
        print(f"Session started. User ID: {new_id}") # Print User ID
        user_dictionary[new_id] = []
        
        # INSTRUCTIONS, instructions must be in quotations, \n means a new line to the AI (not needed).
        instructions = (
            "You are a Chat Bot for the KU HMC Lab\n"
            "instruction...\n"
            "instruction...\n"
        )
        
        # Add ID and instructions to users session data
        user_dictionary[new_id].append({
            "role": "system",
            "content": instructions
        })
        
        # Create a GitHub repository file to store the Data
        update_github_log(new_id)
            
        return new_id, f"<script>window.parent.postMessage({{user_id:'{new_id}'}}, '*')</script>"

    # User message function so we can show the messages separately
    def add_user_message(user_input, user_id):
        if not user_id:
            return [], ""
        if user_id not in user_dictionary:
            user_dictionary[user_id] = []
            
        # Add messages to users session data
        user_dictionary[user_id].append({
            "role": "user",
            "content": user_input
        })
        
        # Format conversation
        formatted_chat = [
            msg for msg in user_dictionary[user_id]
            if msg["role"] != "system"
        ]
        
        return formatted_chat, ""

    # Optional function to simulate the time it takes a "person" to read the users message, (amount of time reading, bubbles dont show up).
    def pause_before_typing(user_id):
        # Calculate reading time based on the length of the user's message
        if user_id and user_id in user_dictionary and len(user_dictionary[user_id]) > 0:
            user_msg = user_dictionary[user_id][-1]["content"]
            
            # Fast reading speed, caps between 1.0 and 3.0 seconds
            read_delay = max(1.0, min(len(user_msg) / 30, 3.0)) 
            time.sleep(read_delay)
        else:
            time.sleep(1)

    # Generate Chat Bot response
    def predict_prompt(user_id):
        if not user_id or user_id not in user_dictionary:
            return []
            
        # Try to get the response from OpenAI
        try:
            response = client.chat.completions.create(
                model="gpt-4o", # Model Type, must be under quotations, for example model="gpt-4o-mini".
                messages=user_dictionary[user_id]
            )
            reply = response.choices[0].message.content
            
            # Optional delay for a "human" response time based on the length of the chatbot's reply (amount of time typing message, bubbles show up)
            typing_delay = max(1.0, min(len(reply) / 40, 5.0)) # Caps typing delay between 1.0 and 5.0 seconds
            time.sleep(typing_delay)

        # If we didnt get the response find and show the Error
        except Exception as e:
            error_msg = str(e).lower()
            # Catch invalid or expired API keys and show a error message
            if "invalid_api_key" in error_msg or "incorrect api key" in error_msg:
                gr.Warning("The OpenAI API Key provided is invalid or expired.")
                reply = "System Error: The OpenAI API Key is invalid."
            elif "ascii" in error_msg:
                gr.Warning("OPENAI_KEY was set incorrectly.")
                reply = "System Error: The OpenAI API Key is improperly formatted."
            else:
                gr.Warning(f"OpenAI Error: {str(e)}")
                reply = f"System Error: {str(e)}"
            
        # Add chatbot reply to session data
        user_dictionary[user_id].append({
            "role": "assistant",
            "content": reply
        })
        
        # Format conversation
        formatted_chat = [
            msg for msg in user_dictionary[user_id]
            if msg["role"] != "system"
        ]
        
        # Save to GitHub
        update_github_log(user_id)
            
        return formatted_chat

    # User Interface components
    Chatbot = gr.Chatbot(
        type="messages",
        label="Anonymous User" # Participant User Name
    )

    # User textbox for them to type in
    txt = gr.Textbox(
        show_label=False,
        placeholder="Please write your message and press Enter",
        container=False
    )
    
    # Submits
    txt.submit(
        add_user_message,
        inputs=[txt, user_id],
        outputs=[Chatbot, txt],
        queue=False
    ).then(
        pause_before_typing,
        inputs=[user_id],
        outputs=None
    ).then(
        predict_prompt,
        inputs=[user_id],
        outputs=Chatbot
    )

# Lanch!
chatblock.launch()