from pydantic import BaseModel, Field
from restack_ai.workflow import workflow, import_functions, log
from src.functions.llm_chat import llm_chat, LlmChatInput, Message


class AgentChatInput(BaseModel):
    message: str


class MessageEvent(BaseModel):
    content: str


class EndEvent(BaseModel):
    end: bool


@workflow.defn()
class AgentChat:
    def __init__(self) -> None:
        self.end = False
        self.messages = []
        
    @workflow.event
    async def message(self, message: MessageEvent) -> List[Message]:
        self.messages.append({"role": "user", "content": message.content})
        assistant_message = await workflow.step(llm_chat, LlmChatInput(messages=self.messages))
        self.messages.append(assistant_message)
        return assistant_message
      
    @workflow.event
    async def end(self, end: EndEvent) -> EndEvent:
        self.end = True
        return end
      
    @workflow.run
    async def run(self, input: AgentChatInput):
        await workflow.step(llm_chat, LlmChatInput(messages=self.messages))
        await workflow.condition(lambda: self.end)
        return