from datetime import datetime

from dash import (
    html, 
    Dash
)
from dash.dependencies import (
    Input, 
    Output, 
    State
)

from llm_utils.client import ChatGPT


def format_chat_messages(chat_history):
    chat_messages = []
    for message in chat_history:
        chat_messages.append(html.Div([
            html.P(f'{message["sender"]}: {message["message"]}'),
            html.P(f'Sent at: {message["timestamp"]}')
        ]))
    return chat_messages


def register_callbacks(app: Dash):
    
    chat_gpt = ChatGPT(model="gpt4")
    
    @app.callback(
        [Output('chat-container', 'children'),
        Output('chat-history', 'data')],
        [Input('send-button', 'n_clicks')],
        [State('user-input', 'value'),
        State('chat-history', 'data')]
    )
    def update_chat(n_clicks, input_value, chat_history):
        if chat_history is None:
            chat_history = []
        
        if n_clicks > 0 and input_value:
            chat_history.append({
                'sender': 'User',
                'message': input_value,
                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })
            
            response = chat_gpt.chat_with_gpt(input_value)
            
            # Add response to chat history
            chat_history.append({
                'sender': 'Language Model',
                'message': response,
                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })
        
        return format_chat_messages(chat_history), chat_history