import logging
import re
import time
from collections.abc import Generator
from threading import Thread
from typing import Union, cast

from sqlalchemy import select
from sqlalchemy.orm import Session

from constants.tts_auto_play_timeout import TTS_AUTO_PLAY_TIMEOUT, TTS_AUTO_PLAY_YIELD_CPU_TIME
from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
from core.app.entities.app_invoke_entities import (
    AgentChatAppGenerateEntity,
    ChatAppGenerateEntity,
    CompletionAppGenerateEntity,
)
from core.app.entities.queue_entities import (
    QueueAgentMessageEvent,
    QueueAgentThoughtEvent,
    QueueAnnotationReplyEvent,
    QueueErrorEvent,
    QueueLLMChunkEvent,
    QueueMessageEndEvent,
    QueueMessageFileEvent,
    QueueMessageReplaceEvent,
    QueuePingEvent,
    QueueRetrieverResourcesEvent,
    QueueStopEvent,
)
from core.app.entities.task_entities import (
    AgentMessageStreamResponse,
    AgentThoughtStreamResponse,
    ChatbotAppBlockingResponse,
    ChatbotAppStreamResponse,
    CompletionAppBlockingResponse,
    CompletionAppStreamResponse,
    EasyUITaskState,
    ErrorStreamResponse,
    MessageAudioEndStreamResponse,
    MessageAudioStreamResponse,
    MessageEndStreamResponse,
    StreamEvent,
    StreamResponse,
)
from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline
from core.app.task_pipeline.message_cycle_manager import MessageCycleManager
from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk
from core.file import helpers as file_helpers
from core.file.enums import FileTransferMethod
from core.model_manager import ModelInstance
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
from core.model_runtime.entities.message_entities import (
    AssistantPromptMessage,
    TextPromptMessageContent,
)
from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
from core.ops.entities.trace_entity import TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
from core.prompt.utils.prompt_message_util import PromptMessageUtil
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from core.tools.signature import sign_tool_file
from events.message_event import message_was_created
from extensions.ext_database import db
from libs.datetime_utils import naive_utc_now
from models.model import (
    AppMode,
    Conversation,
    LLMGenerationDetail,
    Message,
    MessageAgentThought,
    MessageFile,
    UploadFile,
)

logger = logging.getLogger(__name__)


class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline):
    """
    EasyUIBasedGenerateTaskPipeline is a class that generate stream output and state management for Application.
    """

    _THINK_PATTERN = re.compile(r"<think[^>]*>(.*?)</think>", re.IGNORECASE | re.DOTALL)

    _task_state: EasyUITaskState
    _application_generate_entity: Union[ChatAppGenerateEntity, CompletionAppGenerateEntity, AgentChatAppGenerateEntity]
    _precomputed_event_type: StreamEvent | None = None

    def __init__(
        self,
        application_generate_entity: Union[
            ChatAppGenerateEntity, CompletionAppGenerateEntity, AgentChatAppGenerateEntity
        ],
        queue_manager: AppQueueManager,
        conversation: Conversation,
        message: Message,
        stream: bool,
    ):
        super().__init__(
            application_generate_entity=application_generate_entity,
            queue_manager=queue_manager,
            stream=stream,
        )
        self._model_config = application_generate_entity.model_conf
        self._app_config = application_generate_entity.app_config

        self._conversation_id = conversation.id
        self._conversation_mode = conversation.mode

        self._message_id = message.id
        self._message_created_at = int(message.created_at.timestamp())

        self._task_state = EasyUITaskState(
            llm_result=LLMResult(
                model=self._model_config.model,
                prompt_messages=[],
                message=AssistantPromptMessage(content=""),
                usage=LLMUsage.empty_usage(),
            )
        )

        self._message_cycle_manager = MessageCycleManager(
            application_generate_entity=application_generate_entity,
            task_state=self._task_state,
        )

        self._conversation_name_generate_thread: Thread | None = None

    def process(
        self,
    ) -> Union[
        ChatbotAppBlockingResponse,
        CompletionAppBlockingResponse,
        Generator[Union[ChatbotAppStreamResponse, CompletionAppStreamResponse], None, None],
    ]:
        if self._application_generate_entity.app_config.app_mode != AppMode.COMPLETION:
            # start generate conversation name thread
            self._conversation_name_generate_thread = self._message_cycle_manager.generate_conversation_name(
                conversation_id=self._conversation_id, query=self._application_generate_entity.query
            )

        generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager)
        if self.stream:
            return self._to_stream_response(generator)
        else:
            return self._to_blocking_response(generator)

    def _to_blocking_response(
        self, generator: Generator[StreamResponse, None, None]
    ) -> Union[ChatbotAppBlockingResponse, CompletionAppBlockingResponse]:
        """
        Process blocking response.
        :return:
        """
        for stream_response in generator:
            if isinstance(stream_response, ErrorStreamResponse):
                raise stream_response.err
            elif isinstance(stream_response, MessageEndStreamResponse):
                extras = {"usage": self._task_state.llm_result.usage.model_dump()}
                if self._task_state.metadata:
                    extras["metadata"] = self._task_state.metadata.model_dump()
                response: Union[ChatbotAppBlockingResponse, CompletionAppBlockingResponse]
                if self._conversation_mode == AppMode.COMPLETION:
                    response = CompletionAppBlockingResponse(
                        task_id=self._application_generate_entity.task_id,
                        data=CompletionAppBlockingResponse.Data(
                            id=self._message_id,
                            mode=self._conversation_mode,
                            message_id=self._message_id,
                            answer=cast(str, self._task_state.llm_result.message.content),
                            created_at=self._message_created_at,
                            **extras,
                        ),
                    )
                else:
                    response = ChatbotAppBlockingResponse(
                        task_id=self._application_generate_entity.task_id,
                        data=ChatbotAppBlockingResponse.Data(
                            id=self._message_id,
                            mode=self._conversation_mode,
                            conversation_id=self._conversation_id,
                            message_id=self._message_id,
                            answer=cast(str, self._task_state.llm_result.message.content),
                            created_at=self._message_created_at,
                            **extras,
                        ),
                    )

                return response
            else:
                continue

        raise RuntimeError("queue listening stopped unexpectedly.")

    def _to_stream_response(
        self, generator: Generator[StreamResponse, None, None]
    ) -> Generator[Union[ChatbotAppStreamResponse, CompletionAppStreamResponse], None, None]:
        """
        To stream response.
        :return:
        """
        for stream_response in generator:
            if isinstance(self._application_generate_entity, CompletionAppGenerateEntity):
                yield CompletionAppStreamResponse(
                    message_id=self._message_id,
                    created_at=self._message_created_at,
                    stream_response=stream_response,
                )
            else:
                yield ChatbotAppStreamResponse(
                    conversation_id=self._conversation_id,
                    message_id=self._message_id,
                    created_at=self._message_created_at,
                    stream_response=stream_response,
                )

    def _listen_audio_msg(self, publisher: AppGeneratorTTSPublisher | None, task_id: str):
        if publisher is None:
            return None
        audio_msg = publisher.check_and_get_audio()
        if audio_msg and isinstance(audio_msg, AudioTrunk) and audio_msg.status != "finish":
            # audio_str = audio_msg.audio.decode('utf-8', errors='ignore')
            return MessageAudioStreamResponse(audio=audio_msg.audio, task_id=task_id)
        return None

    def _wrapper_process_stream_response(
        self, trace_manager: TraceQueueManager | None = None
    ) -> Generator[StreamResponse, None, None]:
        tenant_id = self._application_generate_entity.app_config.tenant_id
        task_id = self._application_generate_entity.task_id
        publisher = None
        text_to_speech_dict = self._app_config.app_model_config_dict.get("text_to_speech")
        if (
            text_to_speech_dict
            and text_to_speech_dict.get("autoPlay") == "enabled"
            and text_to_speech_dict.get("enabled")
        ):
            publisher = AppGeneratorTTSPublisher(
                tenant_id, text_to_speech_dict.get("voice", None), text_to_speech_dict.get("language", None)
            )
        for response in self._process_stream_response(publisher=publisher, trace_manager=trace_manager):
            while True:
                audio_response = self._listen_audio_msg(publisher, task_id)
                if audio_response:
                    yield audio_response
                else:
                    break
            yield response

        start_listener_time = time.time()
        # timeout
        while (time.time() - start_listener_time) < TTS_AUTO_PLAY_TIMEOUT:
            if publisher is None:
                break
            audio = publisher.check_and_get_audio()
            if audio is None:
                # release cpu
                # sleep 20 ms ( 40ms => 1280 byte audio file,20ms => 640 byte audio file)
                time.sleep(TTS_AUTO_PLAY_YIELD_CPU_TIME)
                continue
            if audio.status == "finish":
                break
            else:
                start_listener_time = time.time()
                yield MessageAudioStreamResponse(audio=audio.audio, task_id=task_id)
        if publisher:
            yield MessageAudioEndStreamResponse(audio="", task_id=task_id)

    def _process_stream_response(
        self, publisher: AppGeneratorTTSPublisher | None, trace_manager: TraceQueueManager | None = None
    ) -> Generator[StreamResponse, None, None]:
        """
        Process stream response.
        :return:
        """
        for message in self.queue_manager.listen():
            if publisher:
                publisher.publish(message)
            event = message.event

            if isinstance(event, QueueErrorEvent):
                with Session(db.engine) as session:
                    err = self.handle_error(event=event, session=session, message_id=self._message_id)
                    session.commit()
                yield self.error_to_stream_response(err)
                break
            elif isinstance(event, QueueStopEvent | QueueMessageEndEvent):
                if isinstance(event, QueueMessageEndEvent):
                    if event.llm_result:
                        self._task_state.llm_result = event.llm_result
                else:
                    self._handle_stop(event)

                # handle output moderation
                output_moderation_answer = self.handle_output_moderation_when_task_finished(
                    cast(str, self._task_state.llm_result.message.content)
                )
                if output_moderation_answer:
                    self._task_state.llm_result.message.content = output_moderation_answer
                    yield self._message_cycle_manager.message_replace_to_stream_response(
                        answer=output_moderation_answer
                    )

                with Session(db.engine) as session:
                    # Save message
                    self._save_message(session=session, trace_manager=trace_manager)
                    session.commit()
                message_end_resp = self._message_end_to_stream_response()
                yield message_end_resp
            elif isinstance(event, QueueRetrieverResourcesEvent):
                self._message_cycle_manager.handle_retriever_resources(event)
            elif isinstance(event, QueueAnnotationReplyEvent):
                annotation = self._message_cycle_manager.handle_annotation_reply(event)
                if annotation:
                    self._task_state.llm_result.message.content = annotation.content
            elif isinstance(event, QueueAgentThoughtEvent):
                agent_thought_response = self._agent_thought_to_stream_response(event)
                if agent_thought_response is not None:
                    yield agent_thought_response
            elif isinstance(event, QueueMessageFileEvent):
                response = self._message_cycle_manager.message_file_to_stream_response(event)
                if response:
                    yield response
            elif isinstance(event, QueueLLMChunkEvent | QueueAgentMessageEvent):
                chunk = event.chunk
                delta_text = chunk.delta.message.content
                if delta_text is None:
                    continue
                if isinstance(chunk.delta.message.content, list):
                    delta_text = ""
                    for content in chunk.delta.message.content:
                        logger.debug(
                            "The content type %s in LLM chunk delta message content.: %r", type(content), content
                        )
                        if isinstance(content, TextPromptMessageContent):
                            delta_text += content.data
                        elif isinstance(content, str):
                            delta_text += content  # failback to str
                        else:
                            logger.warning(
                                "Unsupported content type %s in LLM chunk delta message content.: %r",
                                type(content),
                                content,
                            )
                            continue

                if not self._task_state.llm_result.prompt_messages:
                    self._task_state.llm_result.prompt_messages = chunk.prompt_messages

                # handle output moderation chunk
                should_direct_answer = self._handle_output_moderation_chunk(cast(str, delta_text))
                if should_direct_answer:
                    continue

                current_content = cast(str, self._task_state.llm_result.message.content)
                current_content += cast(str, delta_text)
                self._task_state.llm_result.message.content = current_content

                if isinstance(event, QueueLLMChunkEvent):
                    # Determine the event type once, on first LLM chunk, and reuse for subsequent chunks
                    if not hasattr(self, "_precomputed_event_type") or self._precomputed_event_type is None:
                        self._precomputed_event_type = self._message_cycle_manager.get_message_event_type(
                            message_id=self._message_id
                        )
                    yield self._message_cycle_manager.message_to_stream_response(
                        answer=cast(str, delta_text),
                        message_id=self._message_id,
                        event_type=self._precomputed_event_type,
                    )
                else:
                    yield self._agent_message_to_stream_response(
                        answer=cast(str, delta_text),
                        message_id=self._message_id,
                    )
            elif isinstance(event, QueueMessageReplaceEvent):
                yield self._message_cycle_manager.message_replace_to_stream_response(answer=event.text)
            elif isinstance(event, QueuePingEvent):
                yield self.ping_stream_response()
            else:
                continue
        if publisher:
            publisher.publish(None)
        if self._conversation_name_generate_thread:
            logger.debug("Conversation name generation running as daemon thread")

    def _save_message(self, *, session: Session, trace_manager: TraceQueueManager | None = None):
        """
        Save message.
        :return:
        """
        llm_result = self._task_state.llm_result
        usage = llm_result.usage

        message_stmt = select(Message).where(Message.id == self._message_id)
        message = session.scalar(message_stmt)
        if not message:
            raise ValueError(f"message {self._message_id} not found")
        conversation_stmt = select(Conversation).where(Conversation.id == self._conversation_id)
        conversation = session.scalar(conversation_stmt)
        if not conversation:
            raise ValueError(f"Conversation {self._conversation_id} not found")

        message.message = PromptMessageUtil.prompt_messages_to_prompt_for_saving(
            self._model_config.mode, self._task_state.llm_result.prompt_messages
        )
        message.message_tokens = usage.prompt_tokens
        message.message_unit_price = usage.prompt_unit_price
        message.message_price_unit = usage.prompt_price_unit
        message.answer = (
            PromptTemplateParser.remove_template_variables(cast(str, llm_result.message.content).strip())
            if llm_result.message.content
            else ""
        )
        message.updated_at = naive_utc_now()
        message.answer_tokens = usage.completion_tokens
        message.answer_unit_price = usage.completion_unit_price
        message.answer_price_unit = usage.completion_price_unit
        message.provider_response_latency = time.perf_counter() - self.start_at
        message.total_price = usage.total_price
        message.currency = usage.currency
        self._task_state.llm_result.usage.latency = message.provider_response_latency
        message.message_metadata = self._task_state.metadata.model_dump_json()

        if trace_manager:
            trace_manager.add_trace_task(
                TraceTask(
                    TraceTaskName.MESSAGE_TRACE, conversation_id=self._conversation_id, message_id=self._message_id
                )
            )

        # Save LLM generation detail if there's reasoning_content
        self._save_generation_detail(session=session, message=message, llm_result=llm_result)

        message_was_created.send(
            message,
            application_generate_entity=self._application_generate_entity,
        )

    def _save_generation_detail(self, *, session: Session, message: Message, llm_result: LLMResult) -> None:
        """
        Save LLM generation detail for Completion/Chat/Agent-Chat applications.
        For Agent-Chat, also merges MessageAgentThought records.
        """
        import json

        reasoning_list: list[str] = []
        tool_calls_list: list[dict] = []
        sequence: list[dict] = []
        answer = message.answer or ""

        # Check if this is Agent-Chat mode by looking for agent thoughts
        agent_thoughts = (
            session.query(MessageAgentThought)
            .filter_by(message_id=message.id)
            .order_by(MessageAgentThought.position.asc())
            .all()
        )

        if agent_thoughts:
            # Agent-Chat mode: merge MessageAgentThought records
            content_pos = 0
            cleaned_answer_parts: list[str] = []
            for thought in agent_thoughts:
                # Add thought/reasoning
                if thought.thought:
                    reasoning_text = thought.thought
                    if "<think" in reasoning_text.lower():
                        clean_text, extracted_reasoning = self._split_reasoning_from_answer(reasoning_text)
                        if extracted_reasoning:
                            reasoning_text = extracted_reasoning
                            thought.thought = clean_text or extracted_reasoning
                    reasoning_list.append(reasoning_text)
                    sequence.append({"type": "reasoning", "index": len(reasoning_list) - 1})

                # Add tool calls
                if thought.tool:
                    tool_calls_list.append(
                        {
                            "name": thought.tool,
                            "arguments": thought.tool_input or "",
                            "result": thought.observation or "",
                        }
                    )
                    sequence.append({"type": "tool_call", "index": len(tool_calls_list) - 1})

                # Add answer content if present
                if thought.answer:
                    content_text = thought.answer
                    if "<think" in content_text.lower():
                        clean_answer, extracted_reasoning = self._split_reasoning_from_answer(content_text)
                        if extracted_reasoning:
                            reasoning_list.append(extracted_reasoning)
                            sequence.append({"type": "reasoning", "index": len(reasoning_list) - 1})
                        content_text = clean_answer
                        thought.answer = clean_answer or content_text

                    if content_text:
                        start = content_pos
                        end = content_pos + len(content_text)
                        sequence.append({"type": "content", "start": start, "end": end})
                        content_pos = end
                        cleaned_answer_parts.append(content_text)

            if cleaned_answer_parts:
                merged_answer = "".join(cleaned_answer_parts)
                message.answer = merged_answer
                llm_result.message.content = merged_answer
        else:
            # Completion/Chat mode: use reasoning_content from llm_result
            reasoning_content = llm_result.reasoning_content
            if not reasoning_content and answer:
                # Extract reasoning from <think> blocks and clean the final answer
                clean_answer, reasoning_content = self._split_reasoning_from_answer(answer)
                if reasoning_content:
                    answer = clean_answer
                    llm_result.message.content = clean_answer
                    llm_result.reasoning_content = reasoning_content
                    message.answer = clean_answer
            if reasoning_content:
                reasoning_list = [reasoning_content]
                # Content comes first, then reasoning
                if answer:
                    sequence.append({"type": "content", "start": 0, "end": len(answer)})
                sequence.append({"type": "reasoning", "index": 0})

        # Only save if there's meaningful generation detail
        if not reasoning_list and not tool_calls_list:
            return

        # Check if generation detail already exists
        existing = session.query(LLMGenerationDetail).filter_by(message_id=message.id).first()

        if existing:
            existing.reasoning_content = json.dumps(reasoning_list) if reasoning_list else None
            existing.tool_calls = json.dumps(tool_calls_list) if tool_calls_list else None
            existing.sequence = json.dumps(sequence) if sequence else None
        else:
            generation_detail = LLMGenerationDetail(
                tenant_id=self._application_generate_entity.app_config.tenant_id,
                app_id=self._application_generate_entity.app_config.app_id,
                message_id=message.id,
                reasoning_content=json.dumps(reasoning_list) if reasoning_list else None,
                tool_calls=json.dumps(tool_calls_list) if tool_calls_list else None,
                sequence=json.dumps(sequence) if sequence else None,
            )
            session.add(generation_detail)

    @classmethod
    def _split_reasoning_from_answer(cls, text: str) -> tuple[str, str]:
        """
        Extract reasoning segments from <think> blocks and return (clean_text, reasoning).
        """
        matches = cls._THINK_PATTERN.findall(text)
        reasoning_content = "\n".join(match.strip() for match in matches) if matches else ""

        clean_text = cls._THINK_PATTERN.sub("", text)
        clean_text = re.sub(r"\n\s*\n", "\n\n", clean_text).strip()

        return clean_text, reasoning_content or ""

    def _handle_stop(self, event: QueueStopEvent):
        """
        Handle stop.
        :return:
        """
        model_config = self._model_config
        model = model_config.model

        model_instance = ModelInstance(
            provider_model_bundle=model_config.provider_model_bundle, model=model_config.model
        )

        # calculate num tokens
        prompt_tokens = 0
        if event.stopped_by != QueueStopEvent.StopBy.ANNOTATION_REPLY:
            prompt_tokens = model_instance.get_llm_num_tokens(self._task_state.llm_result.prompt_messages)

        completion_tokens = 0
        if event.stopped_by == QueueStopEvent.StopBy.USER_MANUAL:
            completion_tokens = model_instance.get_llm_num_tokens([self._task_state.llm_result.message])

        credentials = model_config.credentials

        # transform usage
        model_type_instance = model_config.provider_model_bundle.model_type_instance
        model_type_instance = cast(LargeLanguageModel, model_type_instance)
        self._task_state.llm_result.usage = model_type_instance.calc_response_usage(
            model, credentials, prompt_tokens, completion_tokens
        )

    def _message_end_to_stream_response(self) -> MessageEndStreamResponse:
        """
        Message end to stream response.
        :return:
        """
        self._task_state.metadata.usage = self._task_state.llm_result.usage
        metadata_dict = self._task_state.metadata.model_dump()
        return MessageEndStreamResponse(
            task_id=self._application_generate_entity.task_id,
            id=self._message_id,
            metadata=metadata_dict,
        )

    def _record_files(self):
        with Session(db.engine, expire_on_commit=False) as session:
            message_files = session.scalars(select(MessageFile).where(MessageFile.message_id == self._message_id)).all()
            if not message_files:
                return None

            files_list = []
            upload_file_ids = [
                mf.upload_file_id
                for mf in message_files
                if mf.transfer_method == FileTransferMethod.LOCAL_FILE and mf.upload_file_id
            ]
            upload_files_map = {}
            if upload_file_ids:
                upload_files = session.scalars(select(UploadFile).where(UploadFile.id.in_(upload_file_ids))).all()
                upload_files_map = {uf.id: uf for uf in upload_files}

            for message_file in message_files:
                upload_file = None
                if message_file.transfer_method == FileTransferMethod.LOCAL_FILE and message_file.upload_file_id:
                    upload_file = upload_files_map.get(message_file.upload_file_id)

                url = None
                filename = "file"
                mime_type = "application/octet-stream"
                size = 0
                extension = ""

                if message_file.transfer_method == FileTransferMethod.REMOTE_URL:
                    url = message_file.url
                    if message_file.url:
                        filename = message_file.url.split("/")[-1].split("?")[0]  # Remove query params
                elif message_file.transfer_method == FileTransferMethod.LOCAL_FILE:
                    if upload_file:
                        url = file_helpers.get_signed_file_url(upload_file_id=str(upload_file.id))
                        filename = upload_file.name
                        mime_type = upload_file.mime_type or "application/octet-stream"
                        size = upload_file.size or 0
                        extension = f".{upload_file.extension}" if upload_file.extension else ""
                    elif message_file.upload_file_id:
                        # Fallback: generate URL even if upload_file not found
                        url = file_helpers.get_signed_file_url(upload_file_id=str(message_file.upload_file_id))
                elif message_file.transfer_method == FileTransferMethod.TOOL_FILE and message_file.url:
                    # For tool files, use URL directly if it's HTTP, otherwise sign it
                    if message_file.url.startswith("http"):
                        url = message_file.url
                        filename = message_file.url.split("/")[-1].split("?")[0]
                    else:
                        # Extract tool file id and extension from URL
                        url_parts = message_file.url.split("/")
                        if url_parts:
                            file_part = url_parts[-1].split("?")[0]  # Remove query params first
                            # Use rsplit to correctly handle filenames with multiple dots
                            if "." in file_part:
                                tool_file_id, ext = file_part.rsplit(".", 1)
                                extension = f".{ext}"
                            else:
                                tool_file_id = file_part
                                extension = ".bin"
                            url = sign_tool_file(tool_file_id=tool_file_id, extension=extension)
                            filename = file_part

                transfer_method_value = message_file.transfer_method
                remote_url = message_file.url if message_file.transfer_method == FileTransferMethod.REMOTE_URL else ""
                file_dict = {
                    "related_id": message_file.id,
                    "extension": extension,
                    "filename": filename,
                    "size": size,
                    "mime_type": mime_type,
                    "transfer_method": transfer_method_value,
                    "type": message_file.type,
                    "url": url or "",
                    "upload_file_id": message_file.upload_file_id or message_file.id,
                    "remote_url": remote_url,
                }
                files_list.append(file_dict)
            return files_list or None

    def _agent_message_to_stream_response(self, answer: str, message_id: str) -> AgentMessageStreamResponse:
        """
        Agent message to stream response.
        :param answer: answer
        :param message_id: message id
        :return:
        """
        return AgentMessageStreamResponse(
            task_id=self._application_generate_entity.task_id, id=message_id, answer=answer
        )

    def _agent_thought_to_stream_response(self, event: QueueAgentThoughtEvent) -> AgentThoughtStreamResponse | None:
        """
        Agent thought to stream response.
        :param event: agent thought event
        :return:
        """
        with Session(db.engine, expire_on_commit=False) as session:
            agent_thought: MessageAgentThought | None = (
                session.query(MessageAgentThought).where(MessageAgentThought.id == event.agent_thought_id).first()
            )

        if agent_thought:
            return AgentThoughtStreamResponse(
                task_id=self._application_generate_entity.task_id,
                id=agent_thought.id,
                position=agent_thought.position,
                thought=agent_thought.thought,
                observation=agent_thought.observation,
                tool=agent_thought.tool,
                tool_labels=agent_thought.tool_labels,
                tool_input=agent_thought.tool_input,
                message_files=agent_thought.files,
            )

        return None

    def _handle_output_moderation_chunk(self, text: str) -> bool:
        """
        Handle output moderation chunk.
        :param text: text
        :return: True if output moderation should direct output, otherwise False
        """
        if self.output_moderation_handler:
            if self.output_moderation_handler.should_direct_output():
                # stop subscribe new token when output moderation should direct output
                self._task_state.llm_result.message.content = self.output_moderation_handler.get_final_output()
                self.queue_manager.publish(
                    QueueLLMChunkEvent(
                        chunk=LLMResultChunk(
                            model=self._task_state.llm_result.model,
                            prompt_messages=self._task_state.llm_result.prompt_messages,
                            delta=LLMResultChunkDelta(
                                index=0,
                                message=AssistantPromptMessage(content=self._task_state.llm_result.message.content),
                            ),
                        )
                    ),
                    PublishFrom.TASK_PIPELINE,
                )

                self.queue_manager.publish(
                    QueueStopEvent(stopped_by=QueueStopEvent.StopBy.OUTPUT_MODERATION), PublishFrom.TASK_PIPELINE
                )
                return True
            else:
                self.output_moderation_handler.append_new_token(text)

        return False
