import chainlit as cl
from models.chat_state import get_initial_state, add_user_message
from workflows.chat_workflow import chainlit_app

# Chainlit message handler
@cl.on_message
async def process_user_message(message: cl.Message):
    """Process user message and generate response"""
    # Get state from session
    state = cl.user_session.get("state")
    
    # Add user message to state
    state = add_user_message(state, message.content)
    
    # Generate response
    state = await chainlit_app.ainvoke(state)
    
    # Update state in session
    cl.user_session.set("state", state)

@cl.on_chat_start
async def start():
    """
    This function is called when a new chat starts.
    Initialize the chat state and store it in the session
    """  
    # Initialize state
    state = get_initial_state()
    
    # Store state in session
    cl.user_session.set("state", state)



@cl.set_chat_profiles
async def chat_profile():
    return [
        cl.ChatProfile(
            name="mitre_attck_navigator_layer_writer",
            markdown_description="MITRE ATT&CK Navigatorのlayerを生成するチャットボットです。",
            icon="public/icon.jpg",
            starters=[
                cl.Starter(
                    label="Lockbitの攻撃シナリオ",
                    message="Lockbitの攻撃シナリオを生成してください"
                ),
                cl.Starter(
                    label="ランサムウェアギャングが使っているテクニックの頻度",
                    message="最近のランサムウェアギャングが使っているテクニックと頻度を色の濃さで表現してください"
                ),
            ]
        )
    ]