AWPAdapter

@contextmanager
def create(self, app: Runnable, import_string: str) -> Iterator[None]:
    raise NotImplementedError("create")

@classmethod
def create_provider(
    cls,
    fastapi_url: str,
) -> ProviderProtocol:
    raise NotImplementedError("create")

def create_subconversation(self) -> UIBase:
    return self

def create_workflow_ui(self: UIBase, workflow_uuid: str) -> "UI":
    return UI(uibase=self, workflow_uuid=workflow_uuid)

def end_of_thread(self, thread_id: str) -> None:
    thread_info = self._awp_threads.pop(thread_id, None)
    if thread_info:
        thread_info.active = False
        logger.info(f"Ended awp thread: {thread_info}")

def error(
    self,
    workflow_uuid: str,
    sender: Optional[str] = None,
    recipient: Optional[str] = None,
    auto_reply: bool = False,
    uuid: Optional[str] = None,
    short: Optional[str] = None,
    long: Optional[str] = None,
) -> Optional[str]:
    uuid = uuid or str(uuid4().hex)
    return self.process_message(
        Error(
            sender=sender,
            recipient=recipient,
            auto_reply=auto_reply,
            uuid=uuid,
            workflow_uuid=workflow_uuid,
            short=short,
            long=long,
        )
    )

def function_call_execution(
    self,
    workflow_uuid: str,
    sender: Optional[str] = None,
    recipient: Optional[str] = None,
    auto_reply: bool = False,
    uuid: Optional[str] = None,
    function_name: Optional[str] = None,
    call_id: Optional[str] = None,
    retval: Any = None,
) -> Optional[str]:
    uuid = uuid or str(uuid4().hex)
    return self.process_message(
        FunctionCallExecution(
            sender=sender,
            recipient=recipient,
            auto_reply=auto_reply,
            uuid=uuid,
            workflow_uuid=workflow_uuid,
            function_name=function_name,
            call_id=call_id,
            retval=retval,
        )
    )

def get_thread_info_of_awp(self, awp_id: str) -> Optional[AWPThreadInfo]:
    return self._awp_threads.get(awp_id)

def get_thread_info_of_workflow(
    self, workflow_uuid: str
) -> Optional[AWPThreadInfo]:
    thread_info = next(
        (x for x in self._awp_threads.values() if x.workflow_id == workflow_uuid),
        None,
    )
    if thread_info is None:
        logger.error(
            f"Workflow {workflow_uuid} not found in threads: {self._awp_threads}"
        )
        raise RuntimeError(
            f"Workflow {workflow_uuid} not found in threads: {self._awp_threads}"
        )
    return thread_info

def keep_alive(
    self,
    workflow_uuid: str,
    sender: Optional[str] = None,
    recipient: Optional[str] = None,
    auto_reply: bool = False,
    uuid: Optional[str] = None,
) -> Optional[str]:
    uuid = uuid or str(uuid4().hex)
    return self.process_message(
        KeepAlive(
            sender=sender,
            recipient=recipient,
            auto_reply=auto_reply,
            uuid=uuid,
            workflow_uuid=workflow_uuid,
        )
    )

def multiple_choice(
    self,
    workflow_uuid: str,
    sender: Optional[str] = None,
    recipient: Optional[str] = None,
    auto_reply: bool = False,
    uuid: Optional[str] = None,
    prompt: Optional[str] = None,
    choices: Optional[list[str]] = None,
    default: Optional[str] = None,
    single: bool = True,
) -> Optional[str]:
    uuid = uuid or str(uuid4().hex)
    choices = choices or []
    return self.process_message(
        MultipleChoice(
            sender=sender,
            recipient=recipient,
            auto_reply=auto_reply,
            uuid=uuid,
            workflow_uuid=workflow_uuid,
            prompt=prompt,
            choices=choices,
            default=default,
            single=single,
        )
    )

def process_message(self, message: IOMessage) -> Optional[str]:
    try:
        return self.visit(message)
    except Exception as e:
        # log the error and return None
        logger.error(f"Error processing message ({message}): {e}", exc_info=True)
        return None

async def run_thread(
    self, input: RunAgentInput, request: Request
) -> AsyncIterator[str]:
    thread_info = self._awp_threads.get(input.thread_id)
    if thread_info is None:
        logger.error(f"Thread {input.thread_id} not found")
        raise RuntimeError(f"Thread {input.thread_id} not found")

    run_started = RunStartedEvent(
        type=EventType.RUN_STARTED,
        thread_id=thread_info.awp_id,
        run_id=thread_info.run_id,
    )
    yield self._sse_send(run_started, thread_info)

    while not await request.is_disconnected():
        try:
            message = await asyncio.wait_for(
                thread_info.out_queue.get(), timeout=0.5
            )
            yield self._sse_send(message, thread_info)
            if isinstance(message, RunFinishedEvent):
                break
            if isinstance(message, CustomEvent) and message.name == "thread_over":
                run_finished = RunFinishedEvent(
                    type=EventType.RUN_FINISHED,
                    thread_id=thread_info.awp_id,
                    run_id=thread_info.run_id,
                )
                yield self._sse_send(run_finished, thread_info)
                logger.info(f"Thread {input.thread_id} is over")
                self.end_of_thread(input.thread_id)
                break
        except asyncio.TimeoutError:
            await asyncio.sleep(
                0
            )  # Yield control briefly, might not be strictly needed
            continue  # Go back to the top and check if request is still open

    logger.info(f"run thread {input.thread_id} completed")

def send_to_thread(self, thread_id: str, message: str) -> None:
    thread_info = self._awp_threads.get(thread_id)
    if thread_info:
        if not thread_info.active:
            logger.error(f"Thread {thread_id} is not active")
            return
        thread_info.out_queue.put_nowait(message)
    else:
        logger.error(f"Thread {thread_id} not found")

def setup_routes(self) -> APIRouter:
    router = APIRouter()

    @router.post(self.awp_path)
    async def run_agent(
        input: RunAgentInput,
        request: Request,
        user_id: Optional[str] = Depends(self.get_user_id),
    ) -> StreamingResponse:
        headers = {
            "Content-Type": "text/event-stream",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",  # Nginx: prevent buffering
        }

        if input.thread_id in self._awp_threads:
            ## existing thread, resume
            logger.info(f"Resuming thread: {input.thread_id}")
            logger.info(f"Messages: {input.messages}")
            thread_info = self._awp_threads[input.thread_id]
            last_message = input.messages[-1]
            if isinstance(last_message, UserMessage):
                thread_info.input_queue.put_nowait(last_message.content)
            return StreamingResponse(
                self.run_thread(input, request), headers=headers
            )

        ## new thread, create
        workflow_uuid: str = uuid4().hex

        thread_info = AWPThreadInfo(input, workflow_id=workflow_uuid)
        self._awp_threads[input.thread_id] = thread_info
        logger.info(f"Created new thread: {input.thread_id}")

        init_msg = InitiateWorkflowModel(
            user_id=user_id,
            workflow_uuid=workflow_uuid,
            params={},
            name=self.wf_name,
        )

        async def process_messages_in_background(workflow_uuid: str) -> None:
            def a_process_messages_in_background(
                workflow_uuid: str,
            ) -> None:
                workflow_ids.workflow_uuid = workflow_uuid
                self.provider.run(
                    name=init_msg.name,
                    ui=self.create_workflow_ui(workflow_uuid),
                    user_id=user_id if user_id else "None",
                    **init_msg.params,
                )

            await asyncify(a_process_messages_in_background)(workflow_uuid)
            workflow_ids.workflow_uuid = None

        try:
            task = asyncio.create_task(
                process_messages_in_background(workflow_uuid)
            )
            logger.info(f"Started task: {task}")
            # asyncio.create_task(
            #    asyncify(process_messages_in_background)(workflow_uuid)
            # )
        except Exception as e:
            logger.error(f"Error in awp endpoint: {e}", stack_info=True)
        finally:
            ...
            # self.end_of_thread(request.thread_id)
        return StreamingResponse(self.run_thread(input, request), headers=headers)

    @router.get(
        self.discovery_path,
        responses={
            404: {"detail": "Key Not Found"},
            504: {"detail": "Unable to connect to provider"},
        },
    )
    def discovery(
        user_id: Optional[str] = Depends(self.get_user_id),
    ) -> list[WorkflowInfo]:
        try:
            names = self.provider.names
        except FastAgencyConnectionError as e:
            raise HTTPException(status_code=504, detail=str(e)) from e

        try:
            descriptions = [self.provider.get_description(name) for name in names]
        except FastAgencyKeyError as e:
            raise HTTPException(status_code=404, detail=str(e)) from e

        return [
            WorkflowInfo(name=name, description=description)
            for name, description in zip(names, descriptions)
        ]

    return router

def start(
    self,
    *,
    app: "Runnable",
    import_string: str,
    name: Optional[str] = None,
    params: dict[str, Any],
    single_run: bool = False,
) -> None:
    raise NotImplementedError("start")

def suggested_function_call(
    self,
    workflow_uuid: str,
    sender: Optional[str] = None,
    recipient: Optional[str] = None,
    auto_reply: bool = False,
    uuid: Optional[str] = None,
    function_name: Optional[str] = None,
    call_id: Optional[str] = None,
    arguments: Optional[dict[str, Any]] = None,
) -> Optional[str]:
    uuid = uuid or str(uuid4().hex)
    arguments = arguments or {}
    return self.process_message(
        SuggestedFunctionCall(
            sender=sender,
            recipient=recipient,
            auto_reply=auto_reply,
            uuid=uuid,
            workflow_uuid=workflow_uuid,
            function_name=function_name,
            call_id=call_id,
            arguments=arguments,
        )
    )

def system_message(
    self,
    workflow_uuid: str,
    sender: Optional[str] = None,
    recipient: Optional[str] = None,
    auto_reply: bool = False,
    uuid: Optional[str] = None,
    message: Optional[dict[str, Any]] = None,
) -> Optional[str]:
    uuid = uuid or str(uuid4().hex)
    message = message or {}
    return self.process_message(
        SystemMessage(
            sender=sender,
            recipient=recipient,
            auto_reply=auto_reply,
            uuid=uuid,
            workflow_uuid=workflow_uuid,
            message=message,
        )
    )

def text_input(
    self,
    workflow_uuid: str,
    sender: Optional[str] = None,
    recipient: Optional[str] = None,
    auto_reply: bool = False,
    uuid: Optional[str] = None,
    prompt: Optional[str] = None,
    suggestions: Optional[list[str]] = None,
    password: bool = False,
) -> Optional[str]:
    uuid = uuid or str(uuid4().hex)
    suggestions = suggestions or []
    return self.process_message(
        TextInput(
            sender=sender,
            recipient=recipient,
            auto_reply=auto_reply,
            uuid=uuid,
            workflow_uuid=workflow_uuid,
            prompt=prompt,
            suggestions=suggestions,
            password=password,
        )
    )

def text_message(
    self,
    workflow_uuid: str,
    sender: Optional[str] = None,
    recipient: Optional[str] = None,
    auto_reply: bool = False,
    uuid: Optional[str] = None,
    body: Optional[str] = None,
) -> Optional[str]:
    uuid = uuid or str(uuid4().hex)
    return self.process_message(
        TextMessage(
            sender=sender,
            recipient=recipient,
            auto_reply=auto_reply,
            uuid=uuid,
            workflow_uuid=workflow_uuid,
            body=body,
        )
    )

def visit(self, message: IOMessage) -> Optional[str]:
    if self.filter and not self.filter(message):
        logger.info(f"Message filtered out: {message}")
        return None
    # call the super class visit method
    return super().visit(message)

def visit_default(self, message: IOMessage) -> Optional[str]:
    async def a_visit_default(
        self: AWPAdapter, message: IOMessage, workflow_uuid: str
    ) -> Optional[str]:
        logger.info(f"Default Visiting message: {message}")

        return None

    if isinstance(message, IOMessage):
        workflow_uuid = message.workflow_uuid
    else:
        logger.error(f"Message is not an IOMessage: {message}")
        logger.error(f"Message type: {type(message)}")
        workflow_uuid = workflow_ids.workflow_uuid

    return syncify(a_visit_default)(self, message, workflow_uuid)

def visit_error(self, message: Error) -> Optional[str]:
    return self.visit_default(message)

def visit_function_call_execution(
    self, message: FunctionCallExecution
) -> Optional[str]:
    return self.visit_default(message)

def visit_input_request(
    self, message: autogen.events.agent_events.InputRequestEvent
) -> None:
    async def a_visit_input_request(
        self: AWPAdapter,
        message: autogen.events.agent_events.InputRequestEvent,
        workflow_uuid: str,
    ) -> None:
        logger.info(f"Visiting input request: {message}")
        thread_info = self.get_thread_info_of_workflow(workflow_uuid)
        if thread_info is None:
            logger.error(
                f"Thread info not found for workflow {workflow_uuid}: {self._awp_threads}"
            )
            raise KeyError(
                f"Thread info not found for workflow {workflow_uuid}: {self._awp_threads}"
            )

        out_queue = thread_info.out_queue
        uuid = str(uuid4().hex)
        message_started = TextMessageStartEvent(
            type=EventType.TEXT_MESSAGE_START, message_id=uuid, role="assistant"
        )
        out_queue.put_nowait(message_started)

        prompt = message.content.prompt.replace(
            "Press enter to skip and use auto-reply",
            "Answer continue to skip and use auto-reply",
        )

        message_content = TextMessageContentEvent(
            type=EventType.TEXT_MESSAGE_CONTENT,
            message_id=uuid,
            delta=prompt,
        )
        out_queue.put_nowait(message_content)

        message_end = TextMessageEndEvent(
            type=EventType.TEXT_MESSAGE_END, message_id=uuid
        )
        out_queue.put_nowait(message_end)

        ## send end of run message, so that the UI can acquire answer and call us back

        run_finished = RunFinishedEvent(
            type=EventType.RUN_FINISHED,
            thread_id=thread_info.awp_id,
            run_id=thread_info.run_id,
        )
        out_queue.put_nowait(run_finished)
        input_queue = thread_info.input_queue
        response = await input_queue.get()
        if response == "continue":
            response = ""
        message.content.respond(response)

    workflow_uuid = workflow_ids.workflow_uuid
    syncify(a_visit_input_request)(self, message, workflow_uuid)

def visit_keep_alive(self, message: KeepAlive) -> Optional[str]:
    return self.visit_default(message)

def visit_multiple_choice(self, message: MultipleChoice) -> Optional[str]:
    return self.visit_default(message)

def visit_run_completion(
    self, message: autogen.events.agent_events.RunCompletionEvent
) -> None:
    async def a_visit_run_completion(
        self: AWPAdapter,
        message: autogen.events.agent_events.RunCompletionEvent,
        workflow_uuid: str,
    ) -> None:
        logger.info(f"Visiting run completion: {message}")
        thread_info = self.get_thread_info_of_workflow(workflow_uuid)
        if thread_info is None:
            logger.error(
                f"Thread info not found for workflow {workflow_uuid}: {self._awp_threads}"
            )
            return
        out_queue = thread_info.out_queue

        thread_over = CustomEvent(
            type=EventType.CUSTOM, name="thread_over", value={}
        )
        out_queue.put_nowait(thread_over)

    workflow_uuid = workflow_ids.workflow_uuid
    return syncify(a_visit_run_completion)(self, message, workflow_uuid)

def visit_suggested_function_call(
    self, message: SuggestedFunctionCall
) -> Optional[str]:
    return self.visit_default(message)

def visit_system_message(self, message: SystemMessage) -> Optional[str]:
    return self.visit_default(message)

def visit_text(self, message: autogen.messages.agent_messages.TextMessage) -> None:
    async def a_visit_text(
        self: AWPAdapter,
        message: autogen.messages.agent_messages.TextMessage,
        workflow_uuid: str,
    ) -> None:
        logger.info(f"Visiting text event: {message}")
        thread_info = self.get_thread_info_of_workflow(workflow_uuid)
        if thread_info is None:
            logger.error(
                f"Thread info not found for workflow {workflow_uuid}: {self._awp_threads}"
            )
            return

        out_queue = thread_info.out_queue
        content = message.content
        uuid = str(content.uuid)
        if content.content:
            message_started = TextMessageStartEvent(
                type=EventType.TEXT_MESSAGE_START, message_id=uuid, role="assistant"
            )
            out_queue.put_nowait(message_started)

            message_content = TextMessageContentEvent(
                type=EventType.TEXT_MESSAGE_CONTENT,
                message_id=uuid,
                delta=content.content,
            )
            out_queue.put_nowait(message_content)

            message_end = TextMessageEndEvent(
                type=EventType.TEXT_MESSAGE_END, message_id=uuid
            )
            out_queue.put_nowait(message_end)

    workflow_uuid = workflow_ids.workflow_uuid
    syncify(a_visit_text)(self, message, workflow_uuid)

def visit_text_input(self, message: TextInput) -> str:
    async def a_visit_text_input(self: AWPAdapter, message: TextInput) -> str:
        workflow_uuid = message.workflow_uuid
        thread_info = self.get_thread_info_of_workflow(workflow_uuid)
        if thread_info is None:
            logger.error(
                f"Thread info not found for workflow {workflow_uuid}: {self._awp_threads}"
            )
            raise KeyError(
                f"Thread info not found for workflow {workflow_uuid}: {self._awp_threads}"
            )

        out_queue = thread_info.out_queue

        message_started = TextMessageStartEvent(
            type=EventType.TEXT_MESSAGE_START,
            message_id=message.uuid,
            role="assistant",
        )
        out_queue.put_nowait(message_started)

        if message.prompt:
            prompt = message.prompt.replace(
                "Press enter to skip and use auto-reply",
                "Answer continue to skip and use auto-reply",
            )
        message_content = TextMessageContentEvent(
            type=EventType.TEXT_MESSAGE_CONTENT,
            message_id=message.uuid,
            delta=prompt,
        )
        out_queue.put_nowait(message_content)

        message_end = TextMessageEndEvent(
            type=EventType.TEXT_MESSAGE_END, message_id=message.uuid
        )
        out_queue.put_nowait(message_end)

        if thread_info.has_text_input_widget():
            # todo : invoke function to get an answer
            ...

        ## send end of run message, so that the UI can acquire answer and call us back
        run_finished = RunFinishedEvent(
            type=EventType.RUN_FINISHED,
            thread_id=thread_info.awp_id,
            run_id=thread_info.run_id,
        )
        out_queue.put_nowait(run_finished)

        # wait for the answer to be sent back
        response = await thread_info.input_queue.get()
        if response == "continue":
            response = ""
        return response

    return syncify(a_visit_text_input)(self, message)

def visit_text_message(self, message: TextMessage) -> None:
    async def a_visit_text_message(self: AWPAdapter, message: TextMessage) -> None:
        workflow_uuid = message.workflow_uuid
        thread_info = self.get_thread_info_of_workflow(workflow_uuid)
        if thread_info is None:
            logger.error(
                f"Thread info not found for workflow {workflow_uuid}: {self._awp_threads}"
            )
            return
        out_queue = thread_info.out_queue

        message_started = TextMessageStartEvent(
            type=EventType.TEXT_MESSAGE_START,
            message_id=message.uuid,
            role="assistant",
        )
        out_queue.put_nowait(message_started)

        message_content = TextMessageContentEvent(
            type=EventType.TEXT_MESSAGE_CONTENT,
            message_id=message.uuid,
            delta=message.body,
        )
        out_queue.put_nowait(message_content)

        message_end = TextMessageEndEvent(
            type=EventType.TEXT_MESSAGE_END, message_id=message.uuid
        )
        out_queue.put_nowait(message_end)

    syncify(a_visit_text_message)(self, message)

def visit_tool_call(
    self, message: autogen.messages.agent_messages.ToolCallMessage
) -> None:
    async def a_visit_tool_call(
        self: AWPAdapter,
        message: autogen.messages.agent_messages.ToolCallMessage,
        workflow_uuid: str,
    ) -> None:
        logger.info(f"Visiting tool call event: {message}")
        thread_info = self.get_thread_info_of_workflow(workflow_uuid)
        if thread_info is None:
            logger.error(
                f"Thread info not found for workflow {workflow_uuid}: {self._awp_threads}"
            )
            return

        out_queue = thread_info.out_queue
        content = message.content
        uuid = str(content.uuid)
        message_started = TextMessageStartEvent(
            type=EventType.TEXT_MESSAGE_START, message_id=uuid, role="assistant"
        )
        out_queue.put_nowait(message_started)

        message_content = TextMessageContentEvent(
            type=EventType.TEXT_MESSAGE_CONTENT,
            message_id=uuid,
            delta=f"AG2 wants to invoke tool: {content.tool_calls[0].function.name}",
        )
        out_queue.put_nowait(message_content)

        message_end = TextMessageEndEvent(
            type=EventType.TEXT_MESSAGE_END, message_id=uuid
        )
        out_queue.put_nowait(message_end)

    workflow_uuid = workflow_ids.workflow_uuid
    syncify(a_visit_tool_call)(self, message, workflow_uuid)