import asyncio
import datetime
import json
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Optional

import dataiku
from dataikuapi.dss.llm import DSSLLMStreamedCompletionFooter
from dataikuapi.dss.llm_tracing import SpanBuilder
from flask import current_app, g, request

from backend.agents.agent_runtime_builder import build_agent, build_orchestrator
from backend.agents.user_agent import UserAgent
from backend.config import (
    agents_as_tools,
    get_asw_chat_llm_id,
    get_asw_chat_system_prompt,
    get_conversation_vision_llm,
    get_default_llm_id,
    get_doc_generation_enabled,
    get_embed_link_settings,
    get_enable_byo_for_agents,
    get_enable_upload_documents,
    get_enterprise_agents,
    get_extraction_mode,
    get_global_system_prompt,
    get_guardrails_enabled,
    get_guardrails_pattern,
    get_orchestrator_llm,
    get_quota_images_per_conversation,
    get_uploads_managedfolder_id,
    has_orchestrator_llm,
)
from backend.constants import ExtractionMode
from backend.models.events import EventKind
from backend.schemas import schemas
from backend.services.artifact_read_service.service import CURRENT_ASSISTANT_MESSAGE_ID_CONTEXT_KEY
from backend.services.derived_documents_service import (
    count_conversation_images,
    get_structured_documents,
    predict_new_image_count,
    process_derived_documents,
)
from backend.services.generated_document.generated_document_revision_persistence import (
    append_messages_with_generated_document_transition,
)
from backend.services.guardrails_service import emit_guardrails_filter_events, process_documents_for_guardrails
from backend.services.helpers.conversation_agent_resolution import (
    ResolvedConversationAgents,
    filter_active_agents,
    get_agent_name,
    get_user_agent_full,
    load_user_agents_full,
    missing_agents,
    resolve_requested_agents,
)
from backend.services.helpers.history_artifacts import add_history_artifact_summaries
from backend.services.helpers.history_documents import add_history_documents
from backend.services.helpers.non_multimodal_turn_preparation import prepare_turn_for_non_multimodal
from backend.services.helpers.turn_routing import (
    ROUTE_CONTEXT_KEY,
    RouteKind,
    TurnRoute,
    history_for_route,
    resolve_turn_route,
    route_from_context,
)
from backend.services.orchestrator_service import OrchestratorService
from backend.services.stream_link_enrichment_service import (
    StreamLinkEnrichmentService,
    serialize_messages_with_enriched_links,
)
from backend.services.title_generation import (
    generate_title_with_trace,
    update_conversation_title,
)
from backend.utils.connection_credentials_utils import (
    CONNECTION_CREDENTIALS_ACTION_KEY,
    append_connection_credentials_issue_from_response,
    get_connection_credentials_issue,
)
from backend.utils.conv_utils import get_selected_agents_as_objs, normalise_stream_event
from backend.utils.events_utils import (
    get_chart_plans,
    get_references,
    get_selected_agents,
    get_used_agent_ids,
    get_used_tables,
)
from backend.utils.llm_utils import (
    add_completion_msgs,
    get_llm_display_name,
    get_memory_fragment_msg,
    get_tool_validation_requests_msg,
    get_tool_validation_responses_msg,
    is_agent_like,
    normalize_llm_id,
)
from backend.utils.logging_utils import extract_error_message, get_logger, sanitize_for_log
from backend.utils.tool_validation import find_pending_validation_message, pending_validation_ids
from backend.utils.user_utils import get_agent_context
from backend.utils.utils import (
    call_dss_agent_full_conversation,
    select_agents,
)

logger = get_logger(__name__)


class ConversationError(Exception):
    """Raised when payload validation fails or other business rules break."""


def build_actions(events: list[dict]) -> dict:
    """Derive frontend message actions from the assistant event log."""
    from flask import g

    actions = {}
    connection_issue = get_connection_credentials_issue(events)
    if connection_issue:
        actions[CONNECTION_CREDENTIALS_ACTION_KEY] = connection_issue

    if events:
        chart_plans = get_chart_plans(events)
        if chart_plans:
            actions["chart_plans"] = chart_plans
        references = get_references(events)
        sel_agents = get_selected_agents(events)
        logger.info(f"Building message actions, references: {len(references)}")
        logger.info(f"Building message actions, selected agents: {len(sel_agents)}")
        if len(references) != 1 or len(sel_agents) != 1:
            # Only enable stories for one agent and one tool with references
            actions["stories"] = {"enable_stories": False}
        else:
            tables = get_used_tables(references[0])
            logger.info(f"Building message actions, tables: {tables}")
            eas = get_enterprise_agents(g.get("authIdentifier", None))
            agent = next((ea for ea in eas if ea.get("id") == sel_agents[0]["agentId"]), None)
            logger.info(f"Building message actions, agent: {agent}")
            if tables and agent and agent.get("stories_workspace"):
                actions["stories"] = {
                    "enable_stories": True,
                    "tables": tables,
                    "agent_id": sel_agents[0]["agentId"],
                    "stories_workspace": agent.get("stories_workspace"),
                    "agent_short_example_queries": agent.get("agent_short_example_queries", []),
                    "agent_example_queries": agent.get("agent_example_queries", []),
                }
    return actions


class ConversationService:
    def __init__(self, store, draft_mode: bool = False, answers_mode: bool = False):
        """Create a per-turn conversation service without preloading any agent catalog."""
        from backend.database.user_store_protocol import IUserStore

        self.store: IUserStore = store
        self.draft_mode = draft_mode  # True only in test contexts
        self.answers_mode = answers_mode
        self.project = dataiku.api_client().get_default_project()
        self._trace_name = "DKU_AGENT_CHAT_QUERY" if answers_mode else "DKU_AGENT_HUB_QUERY"
        self.app_name = "Agent Chat" if answers_mode else "Agent Hub"

    @staticmethod
    def _blank_conversation(conv_id: str, agent_ids: List[str]) -> schemas.FullConversationRead:
        """Return an in-memory fallback conversation object for draft/error branches."""
        return schemas.FullConversationRead(
            conversation_id=conv_id,
            user_id=str(g.get("authIdentifier")),
            title=None,
            agent_ids=agent_ids,
            llm_id="",
            agents_enabled=1,
            doc_gen_enabled=True,
            status=schemas.StatusEnum.ACTIVE,
            messages=[],
        )

    @staticmethod
    def _get_assistant_message_id() -> str:
        """Preallocate the id of the assistant message that will be persisted for this turn."""
        return str(uuid.uuid4())

    @staticmethod
    def _merge_assistant_context(*contexts: Optional[dict], assistant_message_id: str) -> dict:
        """Merge contexts while forcing the runtime assistant message id for this turn."""
        merged: dict = {}
        for context in contexts:
            if isinstance(context, dict):
                merged.update(context)
        merged[CURRENT_ASSISTANT_MESSAGE_ID_CONTEXT_KEY] = assistant_message_id
        return merged

    def _resolve_stream_agents(self, payload: dict, draft_mode: bool):
        """Resolve active agents and history targets for a streaming turn."""
        selected_llm = payload.get("selectedLLM", "")
        agents_enabled = payload.get("modeAgents", True)
        if self.answers_mode:
            # Agent Chat targets one configured LLM or agent, outside normal agent selection.
            agent_ids = []
            agents_enabled = False
            selected_llm = normalize_llm_id(selected_llm or get_asw_chat_llm_id())
            target_is_agent = is_agent_like(selected_llm)
            history_agent_ids = [selected_llm] if target_is_agent else []
            resolved_agents = ResolvedConversationAgents([], {}, {})
            if target_is_agent:
                resolved_agents = resolve_requested_agents(
                    self.store,
                    history_agent_ids,
                    current_user=g.get("authIdentifier"),
                )
            return agent_ids, agents_enabled, selected_llm, history_agent_ids, resolved_agents

        requested_agent_ids = payload.get("agentIds") or []
        resolved_agents = resolve_requested_agents(
            self.store,
            requested_agent_ids,
            current_user=g.get("authIdentifier"),
        )
        agent_ids = filter_active_agents(
            requested_agent_ids,
            resolved_agents,
            current_user=g.get("authIdentifier"),
            allow_draft=draft_mode,
        )
        logger.info("_filter_active_agents called, %s", agent_ids)
        return agent_ids, agents_enabled, selected_llm, agent_ids, resolved_agents

    def _start_title_generation(
        self,
        conv,
        conv_id: str,
        user_msg_dict,
        user_msg: str,
        payload_attachments,
        emit,
        is_first_user_turn: bool,
    ):
        """Start asynchronous title generation for the first user turn when needed."""
        if not is_first_user_turn or (conv.title and conv.title.strip() not in {"New Chat", "Untitled"}):
            return None, None

        app = current_app._get_current_object()
        title_generation_msg = user_msg_dict

        # if the user uploaded a file without any text we add
        # a user msg 'user uploaded file.extension'
        if not user_msg and payload_attachments:
            first_attachment = payload_attachments[0]
            uploaded_filename = first_attachment.get("document_name")
            # we use copy to not mess with the real history
            title_generation_msg = user_msg_dict.model_copy(update={"content": f"user uploaded : {uploaded_filename}"})

        def _generate_title_worker():
            with app.app_context():
                generated_title, title_trace_payload = generate_title_with_trace([title_generation_msg])
                if generated_title:
                    update_conversation_title(self.store, conv_id, generated_title, emit)
                return title_trace_payload

        title_generation_executor = ThreadPoolExecutor(max_workers=1)
        title_generation_future = title_generation_executor.submit(_generate_title_worker)
        return title_generation_future, title_generation_executor

    def _process_stream_attachments(
        self,
        conv_id: str,
        user_msg_dict,
        payload_attachments,
        emit,
        events: list[dict],
    ) -> None:
        """Extract attachment context and record upload processing events for a turn."""
        # Shared events buffer (persisted in assistant message eventLog)
        # Get configuration parameters
        enable_document_upload = get_enable_upload_documents()
        image_quota = get_quota_images_per_conversation()
        extraction_mode = get_extraction_mode()

        # Only process documents if document upload is enabled
        if not payload_attachments or not enable_document_upload:
            return

        # Count existing images in the conversation
        current_image_count = count_conversation_images(self.store, conv_id)

        # Initialize quota_exceeded flag (will be set to True if quota is exceeded)
        quota_exceeded = False

        # Predict how many new images would be generated from attachments (only if screenshots mode)
        predicted_new_images = 0
        if extraction_mode == ExtractionMode.PAGES_SCREENSHOTS.value:
            folder_id = get_uploads_managedfolder_id()
            predicted_new_images = predict_new_image_count(payload_attachments, folder_id)
            logger.info(
                "Predicted %s new images from %s attachments",
                predicted_new_images,
                len(payload_attachments),
            )

            # Calculate total predicted image count
            predicted_total = current_image_count + predicted_new_images

            # Check if quota would be exceeded with new documents
            if predicted_total > image_quota:
                extraction_mode = ExtractionMode.PAGES_TEXT.value
                quota_exceeded = True
                logger.warning(
                    "Image quota WOULD BE EXCEEDED for conversation %s. Predicted total: %s > Quota: %s. "
                    "Switching to TEXT-ONLY mode for all documents.",
                    conv_id,
                    predicted_total,
                    image_quota,
                )
                # Emit event to notify frontend that extraction mode changed to text-only
                evt_mode_change = {
                    "eventKind": EventKind.EXTRACTION_MODE_CHANGED,
                    "eventData": {
                        "extractionMode": ExtractionMode.PAGES_TEXT.value,
                        "reason": "quota_exceeded",
                    },
                }
                emit("chat_event", {**evt_mode_change, "conversationId": conv_id})
                events.append(evt_mode_change)
            else:
                logger.info(
                    "Document extraction mode: %s for conversation %s. Current images: %s/%s, "
                    "Predicted new: %s, Predicted total: %s/%s",
                    extraction_mode,
                    conv_id,
                    current_image_count,
                    image_quota,
                    predicted_new_images,
                    predicted_total,
                    image_quota,
                )

        # Generate derived context (screenshots or text only) based on extraction mode
        # Emit document analysis event (live) and persist it in events buffer
        evt_da = {
            "eventKind": EventKind.DOCUMENT_ANALYSIS,
            "eventData": {"count": len(payload_attachments)},
        }
        emit("chat_event", {**evt_da, "conversationId": conv_id})
        events.append(evt_da)
        try:
            process_derived_documents(
                self.store,
                conv_id,
                g.get("authIdentifier") or "",
                payload_attachments,
                extraction_mode=extraction_mode,
            )
            # Store extraction_mode as JSON with quota_exceeded flag
            self.store.insert_or_update_message_attachments(
                user_msg_dict.id,
                json.dumps(payload_attachments),
                extraction_mode=extraction_mode,
                quota_exceeded=quota_exceeded if extraction_mode == ExtractionMode.PAGES_TEXT.value else False,
            )
            # Mark all attachments as ready after successful processing
            for att in payload_attachments:
                att["uploadStatus"] = "ready"

            # Log the actual new image count
            new_image_count = count_conversation_images(self.store, conv_id)
            images_added = new_image_count - current_image_count

            if extraction_mode == ExtractionMode.PAGES_SCREENSHOTS.value:
                logger.info(
                    "Added %s images to conversation %s. Total: %s/%s (predicted: %s, actual: %s)",
                    images_added,
                    conv_id,
                    new_image_count,
                    image_quota,
                    predicted_new_images,
                    images_added,
                )
            else:
                logger.info(
                    "Processed %s documents in TEXT-ONLY mode. Total images unchanged: %s/%s",
                    len(payload_attachments),
                    new_image_count,
                    image_quota,
                )

            evt_da = {
                "eventKind": EventKind.DOCUMENT_ANALYSIS_COMPLETED,
                "eventData": {"conv_id": conv_id, "documents": payload_attachments},
            }
            emit("chat_event", {**evt_da, "conversationId": conv_id})
            events.append(evt_da)
        except Exception:
            logger.exception("Failed to process derived documents for conversation %s", conv_id)

    # TODO this is redundant with add_history_to_completion in llm_utils.py
    def _build_history(
        self,
        agent_ids: list[str],
        conv_messages: list[dict],
        *,
        documents_map: dict[str, dict] | None = None,
        agents_enabled: bool = True,
        enterprise_agents_by_id: dict[str, dict] | None = None,
    ) -> list[dict]:
        """Build the normalized history payload for the next model or agent call.

        The history includes structured document context when uploads are enabled,
        skips internal tool-validation messages, and prepends the system message
        that matches the current enterprise-agent or base-LLM flow.
        """
        history: list[dict] = []
        docs_map = documents_map or {}
        docs_allowed = (not agents_enabled and get_enable_upload_documents()) or (
            agents_enabled and get_enable_byo_for_agents()
        )
        for _msg in conv_messages:
            if isinstance(_msg, dict):
                msg = _msg
            else:
                msg = _msg.model_dump(mode="json")
            if (
                msg.get("tool_validation_requests")
                or msg.get("tool_validation_responses")
                or msg.get("memory_fragment")
            ):
                if msg.get("tool_validation_requests") or msg.get("tool_validation_responses"):
                    # skip tool validation messages
                    continue
                else:
                    # This is a memory fragment message, we want to send it back as part of the history
                    # Add memory fragment first
                    history.append(get_memory_fragment_msg(msg.get("memory_fragment", {})))
            entry = {
                "role": msg.get("role"),
                "content": msg.get("content", ""),
            }
            add_history_artifact_summaries(entry, msg)
            add_history_documents(entry, msg, docs_map, documents_allowed=docs_allowed)
            history.append(entry)
        if len(agent_ids) == 1:
            agent_id = agent_ids[0]
            enterprise_agents = enterprise_agents_by_id or {}
            enterprise_agent = enterprise_agents.get(agent_id)
            if enterprise_agent is None and enterprise_agents_by_id is None:
                enterprise_agent = next(
                    (ea for ea in get_enterprise_agents(g.get("authIdentifier", None)) if ea.get("id") == agent_id),
                    None,
                )
            if enterprise_agent:
                # Enterprise agent: use agent_system_instructions
                agent_system_instructions = enterprise_agent.get("agent_system_instructions", "")
                today = datetime.date.today().strftime("%Y-%m-%d")
                sys_content = f"Today is {today}"
                if agent_system_instructions:
                    sys_content = f"{sys_content}. {agent_system_instructions}"
                sys_msg = {"role": "system", "content": sys_content}
                return [sys_msg] + list(history)
        today = datetime.date.today().strftime("%Y-%m-%d")

        sys_prompt = get_asw_chat_system_prompt() if self.answers_mode else get_global_system_prompt()
        sys_msg = {"role": "system", "content": f"Today is {today}.{sys_prompt}"}
        return [sys_msg] + list(history)

    def handle_tool_validation_response(self, payload: dict, emit, cancel_event: threading.Event):
        """Resume a paused turn after the user has answered tool validation requests."""
        logger.info("handle_tool_validation_response with Payload: %r", payload)
        conv_id = payload.get("conversationId")
        message_id = payload.get("messageId")
        draft_mode = bool(payload.get("draftMode"))
        client_ip = payload.get("clientIP")
        user_agent = payload.get("userAgent")
        if not conv_id:
            emit("chat_error", {"error": "Missing conversationId"})
            return
        from dataiku.llm.tracing import new_trace

        trace = new_trace(self._trace_name)
        trace.begin(int(time.time() * 1000))

        # -- draft-mode path (legacy full-rewrite) ----------------------
        if draft_mode:
            agent_ids = payload.get("agentIds") or []
            if not agent_ids:
                emit("chat_error", {"error": "Missing agentIds for draft conversation"})
                return
            resolved_agents = resolve_requested_agents(
                self.store,
                agent_ids,
                current_user=g.get("authIdentifier"),
            )
            # Draft chat goes straight to execution, so summaries are no longer enough here:
            # downstream execution and guardrails need the full user-agent payload.
            load_user_agents_full(self.store, resolved_agents, agent_ids)
            return self._stream_draft(
                conv_id,
                agent_ids,
                "",
                emit,
                cancel_event,
                trace,
                resolved_agents,
                payload,
                client_ip=client_ip,
                user_agent=user_agent,
            )

        conv = self.store.get_conversation(conv_id)
        if not conv:
            emit("chat_error", {"error": f"Conversation {conv_id} not found"})
            return

        last_msg = find_pending_validation_message(conv.messages, message_id)
        if not last_msg or not last_msg.tool_validation_requests:
            emit("chat_error", {"error": "Pending tool validation message not found"})
            return

        selected_llm = last_msg.llm_id or conv.llm_id or payload.get("selectedLLM", "")
        if self.answers_mode:
            selected_llm = normalize_llm_id(selected_llm or get_asw_chat_llm_id())
            target_id = selected_llm
            fallback_agent_ids = []
        else:
            target_id = None
            fallback_agent_ids = list(last_msg.selected_agent_ids or [])

        # In Answers, we may still be targeting an agent/augmented LLM; route accordingly.
        route = route_from_context(last_msg.merged_context) or resolve_turn_route(
            answers_mode=self.answers_mode,
            draft_mode=False,
            doc_gen_enabled=get_doc_generation_enabled() and conv.doc_gen_enabled,
            target_id=target_id,
            target_is_agent=is_agent_like(target_id),
            agent_ids=fallback_agent_ids,
            agents_as_tools=agents_as_tools(),
            base_model_id=normalize_llm_id(selected_llm if not fallback_agent_ids else get_default_llm_id()),
            orchestrator_model_id=get_orchestrator_llm(),
        )
        agent_ids = list(route.agent_ids)

        resolved_agents = resolve_requested_agents(
            self.store,
            agent_ids,
            current_user=g.get("authIdentifier"),
        )
        # Agent Chat can target a DSS agent that is not registered in Agent Hub's
        # enterprise/user-agent catalog. This is also allowed on the initial turn.
        if not self.answers_mode:
            if missing := missing_agents(agent_ids, resolved_agents):
                emit("chat_error", {"error": f"Agent(s) not found: {', '.join(missing)}"})
                return
        # Get current responses and requests
        total_requests = last_msg.tool_validation_requests or []
        try:
            # Accumulate responses atomically so retries cannot execute the same validation twice.
            all_tool_responses = self.store.append_tool_validation_response(
                last_msg.id,
                payload.get("toolValidationResponse"),
            )
        except ValueError as exc:
            emit("chat_error", {"error": str(exc)})
            return
        if all_tool_responses is None:
            logger.info("Ignoring duplicate tool validation response for message %s", last_msg.id)
            return

        # If not all responses received, update message and return
        pending_ids = pending_validation_ids(total_requests, all_tool_responses)
        if pending_ids:
            logger.info(
                "Not all tool validation responses received yet (%d/%d), updating message and waiting for more",
                len(total_requests) - len(pending_ids),
                len(total_requests),
            )
            return

        user_login = g.get("authIdentifier") or ""
        documents_map: dict[str, dict] = {}
        try:
            documents_map = get_structured_documents(self.store, conv_id, user_login)
        except Exception:
            logger.exception("Failed to build structured prompt for conversation %s", conv_id)

        # prompt
        history = self._build_history(
            agent_ids,
            conv.messages,
            documents_map=documents_map,
            enterprise_agents_by_id=resolved_agents.enterprise_agents_by_id,
        )
        # Important to keep this order
        history.append(get_memory_fragment_msg(last_msg.memory_fragment or {}))
        history.append(get_tool_validation_requests_msg(last_msg.tool_validation_requests or {}))
        history.append(get_tool_validation_responses_msg(all_tool_responses))

        events: list[dict] = []
        if self.answers_mode:
            target_id = selected_llm
            sel_agents = [
                {
                    "agentId": target_id,
                    "query": last_msg.content,
                    "agentName": get_llm_display_name(target_id),
                }
            ]
            sel_ids = [target_id]
        else:
            sel_ids = agent_ids
            sel_agents = [
                {
                    "agentId": agent_id,
                    "query": last_msg.content,
                    "agentName": get_agent_name(agent_id, resolved_agents=resolved_agents),
                }
                for agent_id in sel_ids
            ]
        assistant_message_id = self._get_assistant_message_id()
        # assistant streaming
        final_reply, artifacts_meta, tool_validation_requests, memory_fragment, context_upsert = self._stream_reply(
            conv_id,
            sel_agents,
            resolved_agents,
            history,
            emit,
            events,
            cancel_event,
            trace,
            last_msg.merged_context,
            assistant_message_id=assistant_message_id,
            client_ip=client_ip,
            user_agent=user_agent,
            doc_gen_enabled=route.doc_gen_enabled,
            route=route,
        )
        append_connection_credentials_issue_from_response(events, trace)

        actions = build_actions(events)
        used_ids = get_used_agent_ids(events)
        assistant_msg = schemas.MessageCreate(
            id=assistant_message_id,
            conversation_id=conv_id,
            role="assistant",
            content=final_reply,
            event_log=events,
            actions=actions,
            artifacts=artifacts_meta,
            tool_validation_requests=tool_validation_requests,
            memory_fragment=memory_fragment,
            selected_agent_ids=sel_ids,
            used_agent_ids=used_ids,
            llm_id=selected_llm,
            agents_enabled=False if self.answers_mode else conv.agents_enabled,
            merged_context=self._merge_assistant_context(
                last_msg.merged_context,
                context_upsert,
                assistant_message_id=assistant_message_id,
            ),
        )
        # atomic write
        append_messages_with_generated_document_transition(
            store=self.store,
            conversation_id=conv_id,
            existing_messages=list(conv.messages or []),
            messages=[assistant_msg],
            assistant_message=assistant_msg,
        )
        # End tracing
        trace.end(int(time.time() * 1000))

        # If we have an assistant message, update it with the trace info
        if assistant_msg is not None:
            assistant_msg.trace = trace.to_dict() if trace else {}
            self.store.update_message(assistant_msg.id, {"trace": assistant_msg.trace})

        # Emit chat_end to ensure frontend stops waiting, even if errors occurred
        if conv_id and not cancel_event.is_set() and not draft_mode:
            fresh_conv = self.store.get_conversation(conv_id) or conv
            title = fresh_conv.title if fresh_conv else "New Conversation"
            messages = fresh_conv.messages if fresh_conv else []

            emit(
                "chat_end",
                {
                    "agentIds": [] if self.answers_mode else (agent_ids or []),
                    "conversationId": conv_id,
                    "title": title,
                    "messages": serialize_messages_with_enriched_links(messages),
                    "modeAgents": False if self.answers_mode else True,
                    "selectedLLM": selected_llm or "",
                    "documentGenerationEnabled": getattr(fresh_conv, "doc_gen_enabled", True),
                    "hasEventLog": bool(events),
                },
            )

    # ------------------------------------------------------------------ #
    #  Streaming send  (WebSocket)
    # ------------------------------------------------------------------ #
    def stream_message(self, payload: dict, emit, cancel_event: threading.Event):
        """Handle the main websocket chat turn from validation through final streaming."""
        logger.info("stream_message called from %s with Payload: %r", request.remote_addr, payload)
        from dataiku.llm.tracing import new_trace

        trace = new_trace(self._trace_name)
        trace.begin(int(time.time() * 1000))

        def emit_draft_chat_end_fallback():
            draft_conv = None
            if agent_ids:
                try:
                    draft_conv_obj = self.store.get_draft_conversation(agent_ids[0])
                    if draft_conv_obj:
                        draft_conv = draft_conv_obj.convo
                except Exception:
                    logger.exception("Failed to load draft conversation for chat_end fallback")
            if not draft_conv:
                draft_conv = self._blank_conversation(conv_id, agent_ids or []).model_dump(mode="json")
            emit(
                "chat_end",
                {
                    "agentIds": agent_ids or [],
                    "conversationId": conv_id,
                    "title": draft_conv.get("title"),
                    "messages": draft_conv.get("messages", []),
                    "hasEventLog": bool(events),
                },
            )

        assistant_msg_dict = None
        conv = {}
        events: list[dict] = []
        agent_ids = None
        resolved_agents = ResolvedConversationAgents([], {}, {})
        draft_mode = False
        conv_id = None
        agents_enabled = None
        selected_llm = ""
        chat_end_emitted = False
        title_generation_future = None
        title_generation_executor = None
        if payload.get("toolValidationResponse"):
            return self.handle_tool_validation_response(payload, emit, cancel_event)
        _t_start = time.perf_counter()
        _t_validation_ms = _t_db_ensure_ms = _t_doc_proc_ms = _t_agent_sel_ms = _t_stream_reply_ms = _t_db_write_ms = (
            _t_first_token_ms
        ) = 0.0
        try:
            draft_mode = bool(payload.get("draftMode"))
            conv_id = payload.get("conversationId")
            user_msg = (payload.get("userMessage") or "").strip()
            doc_gen_enabled = (
                get_doc_generation_enabled() and payload.get("documentGenerationEnabled", True) is not False
            )
            payload_attachments = payload.get("attachments") or []
            client_ip = payload.get("clientIP")
            user_agent = payload.get("userAgent")
            agent_ids, agents_enabled, selected_llm, history_agent_ids, resolved_agents = self._resolve_stream_agents(
                payload, draft_mode
            )
            if self.answers_mode and is_agent_like(selected_llm) and not has_orchestrator_llm():
                # An agent needs the tools orchestrator to retain document-generation capabilities.
                doc_gen_enabled = False
            # -- validation -------------------------------------------------
            _ph = time.perf_counter()
            # Allow empty user message if attachments are provided
            if not conv_id or (not user_msg and not payload_attachments):
                emit("chat_error", {"error": "Missing conversationId or userMessage"})
                return
            if not self.answers_mode:
                miss = missing_agents(agent_ids, resolved_agents)
                if miss:
                    emit("chat_error", {"error": f"Agent(s) not found: {', '.join(miss)}"})
                    return
            _t_validation_ms = (time.perf_counter() - _ph) * 1000

            # -- draft-mode path (legacy full-rewrite) ----------------------
            if draft_mode:
                load_user_agents_full(self.store, resolved_agents, agent_ids)
                self._stream_draft(
                    conv_id,
                    agent_ids,
                    user_msg,
                    emit,
                    cancel_event,
                    trace,
                    resolved_agents,
                    client_ip=client_ip,
                    user_agent=user_agent,
                )
                chat_end_emitted = True
                return

            # -------- non-draft incremental path ---------------------------
            _ph = time.perf_counter()
            conv = self.store.ensure_conversation_exists(
                conv_id,
                agent_ids,
                agents_enabled,
                selected_llm,
                doc_gen_enabled,
            )
            _t_db_ensure_ms = (time.perf_counter() - _ph) * 1000
            old_agent_ids = list(conv.agent_ids or [])
            is_first_user_turn = len(conv.messages) == 0

            # user msg
            user_msg_dict = schemas.MessageCreate(
                id=str(uuid.uuid4()),
                conversation_id=conv_id,
                role="user",
                content=user_msg,
                event_log=[],
                agents_enabled=agents_enabled,
                llm_id=selected_llm,
            )
            # attachments provided by the client (fresh upload metadata),
            if payload_attachments:
                user_msg_dict.attachments = payload_attachments
                # Store attachments in message_attachments table
                self.store.insert_or_update_message_attachments(user_msg_dict.id, json.dumps(payload_attachments))
                # Mark these attachments as 'attached' in derived_documents (if relevant)
                docs = self.store.get_derived_documents(conv_id)
                payload_doc_names = {a["document_name"] for a in payload_attachments if "document_name" in a}
                for doc in docs:
                    if (
                        doc.document_name in payload_doc_names
                        and doc.document_metadata
                        and doc.document_metadata.get("status") == "uploaded"
                    ):
                        doc.document_metadata["status"] = "processed"
                        self.store.upsert_derived_document(
                            conv_id, doc.document_name, doc.document_path, doc.document_metadata
                        )

            conv.messages.append(user_msg_dict)

            # Generate title and title trace
            title_generation_future, title_generation_executor = self._start_title_generation(
                conv,
                conv_id,
                user_msg_dict,
                user_msg,
                payload_attachments,
                emit,
                is_first_user_turn,
            )

            # -------- Document Extraction and Quota Tracking --------
            _ph = time.perf_counter()
            self._process_stream_attachments(conv_id, user_msg_dict, payload_attachments, emit, events)

            _t_doc_proc_ms = (time.perf_counter() - _ph) * 1000

            self._process_conversation_guardrails(conv_id, emit, events)

            # Select based on settings
            _ph = time.perf_counter()
            if self.answers_mode:
                base_llm_id = normalize_llm_id(selected_llm or get_asw_chat_llm_id())
                base_llm_name = get_llm_display_name(base_llm_id)
                sel_agents = [
                    {
                        "agentId": base_llm_id,
                        "query": user_msg,
                        "agentName": base_llm_name,
                    }
                ]
            else:
                sel_agents, justification = select_agents(
                    agent_ids,
                    conv.messages,
                    enterprise_agents_by_id=resolved_agents.enterprise_agents_by_id,
                    user_agents_by_id=resolved_agents.user_agent_summaries_by_id,
                    trace=trace,
                )
                if justification and not agents_as_tools():
                    evt = {
                        "eventKind": EventKind.AGENT_SELECTION,
                        "eventData": {"justification": justification, "selection": sel_agents},
                    }
                    emit("chat_event", evt)
                    events.append(evt)
                elif len(sel_agents):
                    evt = {
                        "eventKind": EventKind.AGENT_SELECTION,
                        "eventData": {"justification": "User selected", "selection": sel_agents},
                    }
                    # emit("chat_event", evt)
                    events.append(evt)
            _t_agent_sel_ms = (time.perf_counter() - _ph) * 1000

            user_login = g.get("authIdentifier") or ""

            def emit_non_multimodal_preparation_event(event_payload: dict) -> None:
                emit("chat_event", {**event_payload, "conversationId": conv_id})
                events.append(event_payload)

            turn_prep = prepare_turn_for_non_multimodal(
                store=self.store,
                project=self.project,
                user_agents_by_id=resolved_agents.user_agent_summaries_by_id,
                conv_id=conv_id,
                user_login=user_login,
                agent_ids=history_agent_ids,
                conv_messages=conv.messages,
                sel_agents=sel_agents,
                agents_enabled=bool(agents_enabled),
                build_history=lambda *args, **kwargs: self._build_history(
                    *args,
                    **kwargs,
                    enterprise_agents_by_id=resolved_agents.enterprise_agents_by_id,
                ),
                emit_event_payload=emit_non_multimodal_preparation_event,
            )
            history = turn_prep.history
            if blocking_error := turn_prep.blocking_error:
                self.store.append_messages(conv_id, [user_msg_dict])
                emit("chat_error", {"error": blocking_error, "conversationId": conv_id})
                return

            for event_payload in turn_prep.event_payloads:
                emit("chat_event", {**event_payload, "conversationId": conv_id})
                events.append(event_payload)

            sel_agents = turn_prep.filtered_sel_agents or []
            sel_ids = [a["agentId"] for a in sel_agents]
            history_by_target = turn_prep.history_by_target

            last_merged_context = {}
            for msg in reversed(conv.messages):
                msg_data = msg if isinstance(msg, dict) else msg.model_dump(mode="json")
                if msg_data.get("role") != "assistant":
                    continue
                candidate = msg_data.get("merged_context")
                if isinstance(candidate, dict):
                    last_merged_context = candidate
                break

            # assistant streaming
            _ph = time.perf_counter()
            _first_token_seen = False

            def _timed_emit(event, data=None):
                nonlocal _t_first_token_ms, _first_token_seen
                if event == "chat_token" and not _first_token_seen:
                    _t_first_token_ms = (time.perf_counter() - _ph) * 1000
                    _first_token_seen = True
                emit(event, data)

            assistant_message_id = self._get_assistant_message_id()
            base_model_id = get_default_llm_id()
            if not self.answers_mode and not agents_enabled:
                base_model_id = selected_llm or base_model_id
            route = resolve_turn_route(
                answers_mode=self.answers_mode,
                draft_mode=False,
                doc_gen_enabled=doc_gen_enabled,
                target_id=selected_llm if self.answers_mode else None,
                target_is_agent=self.answers_mode and is_agent_like(selected_llm),
                agent_ids=sel_ids if not self.answers_mode else [],
                agents_as_tools=agents_as_tools(),
                base_model_id=normalize_llm_id(base_model_id),
                orchestrator_model_id=get_orchestrator_llm(),
            )
            # The route keeps a selected Agent Chat agent available as an orchestrator tool.
            final_reply, artifacts_meta, tool_validation_requests, memory_fragment, context_upsert = self._stream_reply(
                conv_id,
                sel_agents,
                resolved_agents,
                history,
                _timed_emit,
                events,
                cancel_event,
                trace,
                last_merged_context,
                assistant_message_id=assistant_message_id,
                client_ip=client_ip,
                user_agent=user_agent,
                history_by_target=history_by_target,
                doc_gen_enabled=doc_gen_enabled,
                route=route,
            )
            _t_stream_reply_ms = (time.perf_counter() - _ph) * 1000
            append_connection_credentials_issue_from_response(events, trace)
            # Compute merged context for this turn
            base_context = {}
            try:
                base_context = get_agent_context(
                    g.get("authIdentifier"),
                    conv_id,
                    client_ip=client_ip,
                    user_agent=user_agent,
                )
            except Exception:
                logger.exception("Failed to build base agent context for conversation %s", conv_id)
            merged_context = self._merge_assistant_context(
                last_merged_context,
                base_context,
                context_upsert,
                assistant_message_id=assistant_message_id,
            )
            # TODO maybe could optimize and return these directly instead of recomputing them
            actions = build_actions(events)
            used_ids = get_used_agent_ids(events)
            assistant_msg_dict = schemas.MessageCreate(
                id=assistant_message_id,
                conversation_id=conv_id,
                role="assistant",
                content=final_reply,
                event_log=events,
                actions=actions,
                artifacts=artifacts_meta,
                tool_validation_requests=tool_validation_requests,
                memory_fragment=memory_fragment,
                selected_agent_ids=sel_ids,
                used_agent_ids=used_ids,
                llm_id=selected_llm,
                agents_enabled=agents_enabled,
                merged_context=merged_context or None,
            )

            # Add inputs and outputs to trace
            # Inputs: only the last user message (current query)
            trace.inputs["messages"] = [{"role": "user", "text": user_msg}]
            # Outputs: the generated assistant response
            trace.outputs["text"] = final_reply

            # atomic write
            _ph = time.perf_counter()
            append_messages_with_generated_document_transition(
                store=self.store,
                conversation_id=conv_id,
                existing_messages=list(conv.messages or []),
                messages=[user_msg_dict, assistant_msg_dict],
                assistant_message=assistant_msg_dict,
            )
            conv.messages.append(assistant_msg_dict)
            _t_db_write_ms = (time.perf_counter() - _ph) * 1000

            # Update meta for agent/llm changes immediately (title will be updated async if needed)
            if (
                old_agent_ids != agent_ids
                or conv.agents_enabled != agents_enabled
                or conv.llm_id != selected_llm
                or conv.doc_gen_enabled != doc_gen_enabled
            ):
                self.store.update_conversation_meta(
                    conv_id,
                    agent_ids=agent_ids,
                    agents_enabled=agents_enabled,
                    llm_id=selected_llm,
                    doc_gen_enabled=doc_gen_enabled,
                )
        except Exception as e:
            error_message = extract_error_message(str(e))
            logger.exception("Unhandled error in stream_message: %s", error_message)
            emit("chat_error", {"error": f"Error during message processing: {error_message}"})
        finally:
            if title_generation_future is not None:
                try:
                    # we get the generation title trace from the thread
                    if title_trace_payload := title_generation_future.result():
                        trace.append_trace(title_trace_payload)
                except Exception:
                    logger.exception(f"Failed in getting title generation result for conversation {conv_id}")

            if title_generation_executor is not None:
                title_generation_executor.shutdown(wait=False)

            # End tracing
            trace.end(int(time.time() * 1000))

            # If we have an assistant message, update it with the trace info
            if assistant_msg_dict is not None:
                # Add the trace data to the assistant message dictionary.
                assistant_msg_dict.trace = trace.to_dict() if trace else {}
                # Update the message in the store using the update_message function.
                self.store.update_message(assistant_msg_dict.id, {"trace": assistant_msg_dict.trace})

            # Emit chat_end after trace persistence
            if conv_id and not cancel_event.is_set() and not chat_end_emitted:
                if draft_mode:
                    emit_draft_chat_end_fallback()
                else:
                    # Reload conversation from database to get latest messages with extraction_mode (including quota_exceeded)
                    # This ensures extraction_mode JSON is parsed and included in messages sent to frontend
                    fresh_conv = self.store.get_conversation(conv_id)
                    if not fresh_conv and hasattr(conv, "messages"):
                        fresh_conv = conv
                    # Get safe values for chat_end payload
                    title = fresh_conv.title if fresh_conv else "New Conversation"
                    messages = fresh_conv.messages if fresh_conv else []

                    _t_total_ms = (time.perf_counter() - _t_start) * 1000
                    _timing = {
                        "validation_ms": round(_t_validation_ms, 1),
                        "db_ensure_ms": round(_t_db_ensure_ms, 1),
                        "doc_proc_ms": round(_t_doc_proc_ms, 1),
                        "agent_sel_ms": round(_t_agent_sel_ms, 1),
                        "first_token_ms": round(_t_first_token_ms, 1),
                        "stream_reply_ms": round(_t_stream_reply_ms, 1),
                        "db_write_ms": round(_t_db_write_ms, 1),
                        "total_ms": round(_t_total_ms, 1),
                    }
                    logger.info(
                        "(TIMING) stream_message | convId: %s | validation: %.1fms | db_ensure: %.1fms"
                        " | doc_proc: %.1fms | agent_sel: %.1fms | first_token: %.1fms"
                        " | stream_reply: %.1fms | db_write: %.1fms | total: %.1fms",
                        conv_id,
                        _t_validation_ms,
                        _t_db_ensure_ms,
                        _t_doc_proc_ms,
                        _t_agent_sel_ms,
                        _t_first_token_ms,
                        _t_stream_reply_ms,
                        _t_db_write_ms,
                        _t_total_ms,
                    )
                    emit(
                        "chat_end",
                        {
                            "agentIds": agent_ids or [],
                            "conversationId": conv_id,
                            "title": title,
                            "messages": serialize_messages_with_enriched_links(messages),
                            "modeAgents": agents_enabled if agents_enabled is not None else True,
                            "selectedLLM": selected_llm or "",
                            "documentGenerationEnabled": doc_gen_enabled,
                            "hasEventLog": bool(events),
                            "_timing": _timing,
                        },
                    )

    # ------------------------------------------------------------------ #
    #  Draft-mode helper (legacy logic, no optimisation)
    # ------------------------------------------------------------------ #
    def _stream_draft(
        self,
        conv_id: str,
        agent_ids: list[str],
        user_msg: str,
        emit,
        cancel_event: threading.Event,
        trace: SpanBuilder,
        resolved_agents: ResolvedConversationAgents,
        tool_validation_payload: dict = None,
        *,
        client_ip: str | None = None,
        user_agent: str | None = None,
    ):
        """
        Keeps the original behaviour for draft chats: full conversation
        object rewritten on every turn.
        """
        logger.debug("_stream_draft: conv=%s agents=%s user_msg=%r", conv_id, agent_ids, user_msg)
        if tool_validation_payload:
            try:
                conv_data = self.store.append_draft_tool_validation_response(
                    agent_ids[0],
                    tool_validation_payload.get("toolValidationResponse"),
                )
            except ValueError as exc:
                emit("chat_error", {"error": str(exc)})
                return
            if conv_data is None:
                logger.info("Ignoring duplicate tool validation response for draft agent %s", agent_ids[0])
                return
            conv = conv_data
            last_msg = conv_data["messages"][-1] if conv_data["messages"] else None
            if not last_msg or not last_msg.get("toolValidationRequests"):
                emit("chat_error", {"error": "No tool validation requests found in last message"})
                return

            all_tool_responses = last_msg.get("toolValidationResponses", [])
            total_requests = last_msg.get("toolValidationRequests", [])
            pending_ids = pending_validation_ids(total_requests, all_tool_responses)
            if pending_ids:
                logger.info(
                    "Not all tool validation responses received yet (%d/%d), updating message and waiting for more",
                    len(total_requests) - len(pending_ids),
                    len(total_requests),
                )
                return

            history = self._build_history(
                agent_ids,
                conv_data.get("messages", []),
                enterprise_agents_by_id=resolved_agents.enterprise_agents_by_id,
            )
            history.append(get_memory_fragment_msg(last_msg.get("memoryFragment") or {}))
            history.append(get_tool_validation_requests_msg(last_msg.get("toolValidationRequests") or {}))
            history.append(get_tool_validation_responses_msg(all_tool_responses))
            last_user_message = ""
            for msg in reversed(conv_data.get("messages", [])):
                if msg.get("role") == "user":
                    last_user_message = msg.get("content", "")
                    break
            user_msg = last_user_message
        else:
            draft = self.store.get_draft_conversation(agent_ids[0])
            conv = draft.convo if draft else self._blank_conversation(conv_id, agent_ids).model_dump(mode="json")
            user_msg_id = str(uuid.uuid4())
            conv["messages"].append({"id": user_msg_id, "role": "user", "content": user_msg, "hasEventLog": False})

            history = self._build_history(
                agent_ids,
                conv.get("messages", []),
                enterprise_agents_by_id=resolved_agents.enterprise_agents_by_id,
            )

        assistant_message_id = self._get_assistant_message_id()
        events: list[dict] = []
        final_reply, artifacts_meta, tool_validation_requests, memory_fragment, _context_upsert = self._stream_reply(
            conv_id,
            [
                {
                    "agentId": aid,
                    "query": user_msg,
                    "agentName": get_agent_name(
                        aid,
                        resolved_agents=resolved_agents,
                        prefer_full_user_agents=True,
                    ),
                }
                for aid in agent_ids
            ],
            resolved_agents,
            history,
            emit,
            events,
            cancel_event,
            trace,
            # TODO (clement): Currently we don't handle context upsert in draft mode.
            # Probably not a big issue since we don't build structured agents as Quick Agents
            {},
            assistant_message_id=assistant_message_id,
            client_ip=client_ip,
            user_agent=user_agent,
        )
        append_connection_credentials_issue_from_response(events, trace)
        # TODO maybe could optimize and return these directly instead of recomputing them
        actions = build_actions(events)
        used_ids = get_used_agent_ids(events)
        conv["messages"].append(
            {
                "id": assistant_message_id,
                "role": "assistant",
                "content": final_reply,
                "eventLog": events,
                "actions": actions,
                "artifactsMetadata": artifacts_meta,
                "toolValidationRequests": tool_validation_requests,
                "memoryFragment": memory_fragment,
                "selectedAgentIds": agent_ids,
                "usedAgentIds": used_ids,
                "trace": trace.to_dict() if trace else {},
            }
        )
        self.store.upsert_draft_conversation(agent_ids[0], conv)

        if not cancel_event.is_set():
            emit(
                "chat_end",
                {
                    "agentIds": agent_ids,
                    "conversationId": conv_id,
                    "title": conv["title"],
                    "messages": serialize_messages_with_enriched_links(conv["messages"]),
                    "hasEventLog": bool(events),
                },
            )

    # ---------- streaming low-level -----------------------------------
    def _stream_reply(
        self,
        convId,
        sel_agents,
        resolved_agents: ResolvedConversationAgents,
        history,
        emit,
        events,
        cancel_event: threading.Event,
        trace,
        last_merged_context=None,
        assistant_message_id: str | None = None,
        *,
        client_ip: str | None = None,
        user_agent: str | None = None,
        history_by_target: Optional[Dict[str, List[dict]]] = None,
        doc_gen_enabled: bool = True,
        route: TurnRoute | None = None,
    ):
        """Dispatch a prepared turn to the correct runtime path and stream its output."""
        final = ""
        artifacts_meta = {}
        tool_validation_requests = []
        memory_fragment = {}
        generic_error = "Error during agent execution"
        context_upsert = {}
        embed_link_settings = get_embed_link_settings()

        def push_event(ev_dict: dict, store_event: bool = True):
            ev_dict.update({"conversationId": convId})
            if store_event:
                events.append(ev_dict)
            emit("chat_event", ev_dict)

        link_enrichment = StreamLinkEnrichmentService(
            emit_event=push_event,
            cancel_event=cancel_event,
            allowed_origins_enabled=embed_link_settings["allowed_origins_enabled"],
            allowed_origins=embed_link_settings["allowed_origins"],
        )

        def token_cb(tok):
            nonlocal final
            final += tok
            emit("chat_token", {"token": tok, "conversationId": convId})
            link_enrichment.consume_text(tok)

        sel_ids = [a.get("agentId") for a in sel_agents]
        logger.info(f"_stream_reply  called conv: {convId}, sel_ids: {sel_ids}")
        base_context = get_agent_context(
            g.get("authIdentifier"),
            convId,
            client_ip=client_ip,
            user_agent=user_agent,
        )
        merged_context = self._merge_assistant_context(
            last_merged_context,
            base_context,
            assistant_message_id=assistant_message_id or str(uuid.uuid4()),
        )
        merged_context.pop(ROUTE_CONTEXT_KEY, None)
        try:
            target_id = normalize_llm_id(sel_ids[0]) if self.answers_mode and sel_ids else None
            route = route or resolve_turn_route(
                answers_mode=self.answers_mode,
                draft_mode=self.draft_mode,
                doc_gen_enabled=doc_gen_enabled,
                target_id=target_id,
                target_is_agent=is_agent_like(target_id),
                agent_ids=sel_ids if not self.answers_mode else [],
                agents_as_tools=agents_as_tools(),
                base_model_id=get_default_llm_id(),
                orchestrator_model_id=get_orchestrator_llm(),
            )
            logger.info("Using %s route", route.kind.value)

            if route.kind == RouteKind.DIRECT_AGENT:
                target_history = (history_by_target or {}).get(sel_agents[0].get("agentId")) or history
                final, context_upsert = self._stream_single(
                    sel_agents[0],
                    resolved_agents,
                    target_history,
                    merged_context,
                    token_cb,
                    push_event,
                    artifacts_meta,
                    tool_validation_requests,
                    memory_fragment,
                    cancel_event,
                    convId,
                    trace,
                    client_ip=client_ip,
                    user_agent=user_agent,
                )
            elif route.kind == RouteKind.TOOLS_ORCHESTRATOR:
                assert route.model_id
                # Route agent ids are the exact tools exposed to the orchestrator for this turn.
                route_history = history_for_route(route, history, history_by_target, sel_ids)
                final, tool_validation_requests, memory_fragment, context_upsert = self._stream_agent_connect(
                    list(route.agent_ids),
                    resolved_agents,
                    route_history,
                    merged_context,
                    token_cb,
                    push_event,
                    artifacts_meta,
                    cancel_event,
                    convId,
                    trace,
                    base_model=route.model_id,
                    doc_gen_enabled=route.doc_gen_enabled,
                )
            else:
                assert route.model_id
                context = self._get_agent_context_safe(
                    g.get("authIdentifier"),
                    convId,
                    pcb=push_event,
                    agent_name=self.app_name,
                    client_ip=client_ip,
                    user_agent=user_agent,
                )
                if context is None:
                    token_cb(generic_error)
                    return generic_error, artifacts_meta, tool_validation_requests, memory_fragment, context_upsert
                # Selection only needs light metadata, but the orchestrator wrapper still needs full user-agent
                # details to build the concrete DSS agent objects for the selected subset.
                final, tool_validation_requests, memory_fragment = OrchestratorService.stream_multiple_agents(
                    llm_id=route.model_id,
                    sel_agents=sel_agents,
                    messages=history,
                    messages_by_agent_id=history_by_target,
                    context=merged_context,
                    tcb=token_cb,
                    pcb=push_event,
                    artifacts_meta=artifacts_meta,
                    cancel_event=cancel_event,
                    store=self.store,
                    trace=trace,  # Pass trace to orchestrator
                    agents_obj=get_selected_agents_as_objs(
                        self.store,
                        sel_ids,
                        enterprise_agents_by_id=resolved_agents.enterprise_agents_by_id,
                        user_agents=load_user_agents_full(self.store, resolved_agents, sel_ids),
                    ),
                )
            context_upsert[ROUTE_CONTEXT_KEY] = route.to_dict()
            link_enrichment.flush()
            return final, artifacts_meta, tool_validation_requests, memory_fragment, context_upsert
        finally:
            link_enrichment.close()

    # ---- concrete streamers -----------------------------------------
    def _stream_single(
        self,
        agent_info,
        resolved_agents: ResolvedConversationAgents,
        msgs,
        context: Dict,
        tcb,
        pcb,
        artifacts_meta,
        tool_validation_requests,
        memory_fragment,
        cancel_event: threading.Event,
        conv_id,
        trace,
        *,
        client_ip: str | None = None,
        user_agent: str | None = None,
    ):
        """Execute one selected agent, loading full user-agent details only on this path."""
        final = ""
        generic_error = "Error during agent execution"
        final = ""
        context_upsert = {}
        aid = agent_info.get("agentId")
        aname = agent_info.get("agentName")
        user = g.get("authIdentifier")
        agent_query = agent_info.get("query", "")
        if aid in resolved_agents.user_agent_summaries_by_id:
            ua = get_user_agent_full(self.store, resolved_agents, aid)
            if ua is None:
                error_msg = "This agent is no longer accessible."
                tcb(error_msg)
                return error_msg, {}

            logger.info(f"Calling user agent {aid} {aname}")
            pcb({"eventKind": EventKind.CALLING_AGENT, "eventData": {"agentId": aid, "agentName": aname}})
            is_owner = ua.owner == user

            # Same simple logic
            if self.draft_mode and is_owner:
                use_published = False
            else:
                if not ua.published_version:
                    error_msg = "This agent has not been published yet."
                    tcb(error_msg)
                    return error_msg, {}
                use_published = True

            # Check guardrails for agent documents at runtime
            # This re-validates documents in case the guardrails pattern has changed
            guardrails_result = self._enforce_guardrails([aid], pcb, resolved_agents)
            if guardrails_result["blocked_agents"]:
                error_msg = self._emit_guardrails_block_events(guardrails_result["blocked_agents"], pcb, tcb)
                return error_msg, {}

            ua_obj = UserAgent(ua, use_published=use_published)
            from dataiku.llm.tracing import new_trace

            async def _run():
                nonlocal final
                nonlocal context_upsert
                stream_id = str(uuid.uuid4())
                context = self._get_agent_context_safe(
                    user, conv_id, pcb=pcb, agent_id=aid, agent_name=aname, client_ip=client_ip, user_agent=user_agent
                )
                if context is None:
                    tcb(generic_error)
                    final = generic_error
                    return
                # mark usage up-front (we know exactly which agent is called)
                async for ev in ua_obj.aprocess_stream(
                    query={"messages": msgs, "context": context},
                    settings={},
                    trace=trace or new_trace(aid),
                ):
                    if cancel_event.is_set():
                        break
                    footer = ev.get("footer") or {}
                    upsert = footer.get("contextUpsert")
                    if upsert and isinstance(upsert, dict):
                        context_upsert.update(upsert)
                    normalise_stream_event(
                        ev=ev,
                        tcb=tcb,
                        pcb=pcb,
                        msgs=msgs,
                        aid=aid,
                        trace=trace,
                        aname=aname,
                        query=agent_query,
                        artifacts_meta=artifacts_meta,
                        tool_validation_requests=tool_validation_requests,
                        memory_fragment=memory_fragment,
                        stream_id=stream_id,
                    )
                    if "chunk" in ev and "text" in ev["chunk"]:
                        final += ev["chunk"]["text"]

            asyncio.run(_run())
        else:
            logger.info(f"Calling Enterprise Agent: {aid}")
            if not self.answers_mode:
                pcb({"eventKind": EventKind.CALLING_AGENT, "eventData": {"agentId": aid, "agentName": aname}})
            stream_id = str(uuid.uuid4())
            for ev in call_dss_agent_full_conversation(
                aid, aname, msgs, user, True, conv_id=conv_id, trace=trace, context=context
            ):
                if cancel_event.is_set():
                    break
                footer = ev.get("footer") or {}
                upsert = footer.get("contextUpsert")
                if upsert and isinstance(upsert, dict):
                    context_upsert.update(upsert)
                normalise_stream_event(
                    ev=ev,
                    tcb=tcb,
                    pcb=pcb,
                    msgs=msgs,
                    aid=aid,
                    trace=trace,
                    aname=aname,
                    query=agent_query,
                    artifacts_meta=artifacts_meta,
                    tool_validation_requests=tool_validation_requests,
                    memory_fragment=memory_fragment,
                    stream_id=stream_id,
                )
                if "chunk" in ev and "text" in ev["chunk"]:
                    final += ev["chunk"]["text"]
        return final, context_upsert

    def _stream_agent_connect(
        self,
        aids,
        resolved_agents: ResolvedConversationAgents,
        msgs,
        context: Dict,
        tcb,
        pcb,
        artifacts_meta,
        cancel_event: threading.Event,
        conv_id,
        trace,
        *,
        base_model: str,
        doc_gen_enabled: bool = True,
    ):
        """Execute the multi-agent tools path through AgentConnect."""
        final = ""
        tool_validation_reqs = []
        mem_frag = {}
        context_upsert = {}
        from dataiku.llm.tracing import new_trace

        generic_error = "Error during agent execution"
        if not aids:
            logger.info("LLM Mode on")
            agent_connect = build_agent(
                self.store,
                msgs,
                base_model=base_model,
                doc_generation_enabled=doc_gen_enabled,
            )
        elif self.answers_mode:
            # Agent Chat exposes its selected main agent as the only orchestration tool.
            agent_connect = build_orchestrator(
                self.store,
                aids,
                base_model=base_model,
                enterprise_agents_by_id=resolved_agents.enterprise_agents_by_id,
                user_agents=resolved_agents.loaded_user_agents_by_id,
                doc_generation_enabled=doc_gen_enabled,
                agent_chat_mode=True,
            )
        else:
            logger.info(f"Agents Mode on with agents: {len(aids)}")
            # TODO HELA NOT CLEAR TO ME WHAT HAPPENS HERE
            # Check guardrails for all user agents with documents before starting
            # In multi-agent mode, we filter out violating agents instead of blocking entirely
            load_user_agents_full(self.store, resolved_agents, aids)
            guardrails_result = self._enforce_guardrails(aids, pcb, resolved_agents)
            blocked_agent_ids = guardrails_result["blocked_agent_ids"]
            blocked_agents = guardrails_result["blocked_agents"]

            if blocked_agents:
                # Filter out blocked agents
                filtered_aids = [aid for aid in aids if aid not in blocked_agent_ids]

                if not filtered_aids:
                    # All agents are blocked - emit block events and return error
                    error_msg = self._emit_guardrails_block_events(blocked_agents, pcb, tcb)
                    return error_msg, [], {}, {}

                # Some agents remain - emit filter events and continue with remaining agents
                emit_guardrails_filter_events(blocked_agents, pcb)
                aids = filtered_aids

            agent_connect = build_orchestrator(
                self.store,
                aids,
                base_model=base_model,
                enterprise_agents_by_id=resolved_agents.enterprise_agents_by_id,
                user_agents=resolved_agents.loaded_user_agents_by_id,
                doc_generation_enabled=doc_gen_enabled,
            )

        def _run():
            nonlocal final
            nonlocal trace
            if not trace:
                trace = new_trace(self._trace_name)
                trace.begin(int(time.time() * 1000))
            stream_id = str(uuid.uuid4())
            try:
                for ev in agent_connect.process_stream(
                    query={"messages": msgs, "context": context},
                    settings={},
                    artifacts_meta=artifacts_meta,
                    trace=trace,
                    pcb=pcb,
                ):
                    if cancel_event.is_set():
                        break
                    footer = ev.get("footer") or {}
                    upsert = footer.get("contextUpsert")
                    if upsert and isinstance(upsert, dict):
                        context_upsert.update(upsert)
                    normalise_stream_event(
                        ev=ev,
                        tcb=tcb,
                        pcb=pcb,
                        msgs=msgs,
                        stream_id=stream_id,
                        artifacts_meta=artifacts_meta,
                        tool_validation_requests=tool_validation_reqs,
                        memory_fragment=mem_frag,
                        aname=self.app_name,
                    )
                    if "chunk" in ev and "text" in ev["chunk"]:
                        final += ev["chunk"]["text"]
            except Exception as e:
                logger.exception("Error in _stream_agent_connect for conversation %s: %s", conv_id, str(e))
                if not final:
                    tcb(generic_error)
                    final = generic_error
            finally:
                if trace:
                    trace.end(int(time.time() * 1000))

        _run()
        return final, tool_validation_reqs, mem_frag, context_upsert

    # ------------------------------------------------------------------ #
    #  Utility helpers
    # ------------------------------------------------------------------ #
    def _get_agent_context_safe(
        self,
        user: str,
        conv_id: str,
        *,
        pcb=None,
        agent_id: str | None = None,
        agent_name: str | None = None,
        client_ip: str | None = None,
        user_agent: str | None = None,
    ) -> dict | None:
        """Build agent context and surface failures as stream events instead of hard errors."""
        try:
            return get_agent_context(user, conv_id, client_ip=client_ip, user_agent=user_agent)
        except Exception as exc:
            error_message = f"Failed to build agent context: {extract_error_message(str(exc))}"
            logger.exception("Failed to build agent context for conversation %s", conv_id)
            if pcb:
                pcb(
                    {
                        "eventKind": EventKind.AGENT_ERROR,
                        "eventData": {
                            "message": error_message,
                            "agentId": agent_id or "",
                            "agentName": agent_name or self.app_name,
                        },
                    }
                )
            return None

    # TODO (HELA): Move guardrails logic to a separate service class for better separation of concerns. Conv service is too complex already
    def _process_conversation_guardrails(self, conv_id: str, emit, events: list[dict]) -> None:
        """
        Checks all conversation documents against guardrails.
        Emits events for checks, results, and violations.
        """
        if not (get_guardrails_enabled() and get_guardrails_pattern()):
            return

        # Re-check ALL conversation documents (not just new attachments) because:
        # 1. The guardrails pattern may have changed, making previously-failed docs valid
        # 2. The cache in guardrails_service ensures efficiency when patterns haven't changed
        all_conv_docs = self.store.get_derived_documents(conv_id) or []
        all_attachments_for_guardrails = [
            {
                "document_path": doc.document_path,
                "document_name": doc.document_name,
            }
            for doc in all_conv_docs
        ]

        if not all_attachments_for_guardrails:
            return

        try:
            count = len(all_attachments_for_guardrails)
            doc_word = "document" if count == 1 else "documents"
            evt_guardrails_checks = {
                "eventKind": EventKind.GUARDRAILS_CHECKS,
                "eventData": {
                    "message": f"Applying guardrails checks to {count} {doc_word}.",
                },
            }
            emit("chat_event", {**evt_guardrails_checks, "conversationId": conv_id})
            events.append(evt_guardrails_checks)

            guardrails_result = process_documents_for_guardrails(
                self.store,
                conv_id,
                all_attachments_for_guardrails,
            )
            content_violations = guardrails_result.get("content_violations", [])
            extraction_failures = guardrails_result.get("extraction_failures", [])
            all_checked_docs = guardrails_result.get("all_checked", [])

            # Always emit guardrails_checked event with all document statuses for frontend updates
            if all_checked_docs:
                evt_guardrails_checked = {
                    "eventKind": EventKind.GUARDRAILS_CHECKED,
                    "eventData": {
                        "documents": all_checked_docs,
                    },
                }
                emit("chat_event", {**evt_guardrails_checked, "conversationId": conv_id})
                # Don't append to events - this is just for live frontend updates

            # Emit separate events for content violations and extraction failures
            if content_violations:
                violation_names = [d.get("name", "Unknown") for d in content_violations]
                evt_guardrails_violation = {
                    "eventKind": EventKind.GUARDRAILS_VIOLATION,
                    "eventData": {
                        "failed_documents": violation_names,
                        "message": f"Documents filtered due to content restrictions: {', '.join(violation_names)}",
                    },
                }
                emit("chat_event", {**evt_guardrails_violation, "conversationId": conv_id})
                events.append(evt_guardrails_violation)
                logger.warning(
                    f"Guardrails filtered {len(content_violations)} documents for conversation {conv_id}: {violation_names}"
                )

            if extraction_failures:
                extraction_failure_names = [d.get("name", "Unknown") for d in extraction_failures]
                evt_extraction_failed = {
                    "eventKind": EventKind.GUARDRAILS_EXTRACTION_FAILED,
                    "eventData": {
                        "failed_documents": extraction_failure_names,
                        "message": f"Documents excluded due to text extraction failure: {', '.join(extraction_failure_names)}",
                    },
                }
                emit("chat_event", {**evt_extraction_failed, "conversationId": conv_id})
                events.append(evt_extraction_failed)
                logger.warning(
                    f"Text extraction failed for {len(extraction_failures)} documents in conversation {conv_id}: {extraction_failure_names}"
                )
        except Exception:
            logger.exception("Failed to process guardrails for conversation %s", conv_id)

    def _enforce_guardrails(self, agent_ids: list[str], pcb, resolved_agents: ResolvedConversationAgents) -> dict:
        """
        Validates that the specified agents do not violate guardrails.

        Returns a dict with:
            - blocked_agent_ids: set of agent IDs that have violations
            - blocked_agents: list of dicts with agent info and violation details

        The caller is responsible for deciding whether to block or filter,
        and for emitting appropriate events.
        """
        result = {"blocked_agent_ids": set(), "blocked_agents": []}

        if not (get_guardrails_enabled() and get_guardrails_pattern()):
            return result

        user = g.get("authIdentifier")
        load_user_agents_full(self.store, resolved_agents, agent_ids)

        for aid in agent_ids:
            ua = resolved_agents.loaded_user_agents_by_id.get(str(aid))
            # Only user agents have documents to check here; Enterprise agents are skipped or have no ua entry
            if not ua:
                continue

            is_owner = ua.owner == user
            # Standard logic: Use published unless in draft mode AND owner
            use_published = not (self.draft_mode and is_owner)

            docs_to_check = []
            if use_published:
                if ua.published_version:
                    docs_to_check = ua.published_version.documents or []
            else:
                docs_to_check = ua.documents or []

            if not docs_to_check:
                continue

            check_result = self._check_agent_runtime_guardrails(aid, docs_to_check, use_published, pcb)
            if check_result.get("has_violations"):
                content_violations = check_result.get("content_violations", [])
                extraction_failures = check_result.get("extraction_failures", [])
                result["blocked_agent_ids"].add(aid)
                result["blocked_agents"].append(
                    {
                        "agentId": aid,
                        "agentName": ua.name,
                        "content_violations": [d.get("name") for d in content_violations],
                        "extraction_failures": [d.get("name") for d in extraction_failures],
                    }
                )

        return result

    def _emit_guardrails_block_events(self, blocked_agents: list[dict], pcb, tcb=None) -> str:
        """
        Emit guardrails violation/extraction failure events for blocking scenario.
        Returns the combined error message.
        """
        if len(blocked_agents) == 1:
            # Single agent: use single-agent event structure
            ba = blocked_agents[0]
            content_violations = ba["content_violations"]
            extraction_failures = ba["extraction_failures"]
            messages = []

            if content_violations:
                msg = f"Cannot use agent '{ba['agentName']}': some knowledge base documents are blocked by content policy: {', '.join(content_violations)}"
                messages.append(msg)
                evt = {
                    "eventKind": EventKind.GUARDRAILS_VIOLATION,
                    "eventData": {
                        "failed_documents": content_violations,
                        "message": msg,
                        "agentId": ba["agentId"],
                        "agentName": ba["agentName"],
                    },
                }
                pcb(evt)

            if extraction_failures:
                msg = f"Cannot use agent '{ba['agentName']}': text extraction failed for some knowledge base documents: {', '.join(extraction_failures)}"
                messages.append(msg)
                evt = {
                    "eventKind": EventKind.GUARDRAILS_EXTRACTION_FAILED,
                    "eventData": {
                        "failed_documents": extraction_failures,
                        "message": msg,
                        "agentId": ba["agentId"],
                        "agentName": ba["agentName"],
                    },
                }
                pcb(evt)

            combined_msg = " ".join(messages)
            if tcb:
                tcb(combined_msg)
            return combined_msg
        else:
            # Multi-agent event structure
            content_violation_parts = []
            extraction_failure_parts = []

            for ba in blocked_agents:
                agent_name = ba["agentName"]
                content_violations = ba.get("content_violations", [])
                extraction_failures = ba.get("extraction_failures", [])

                if content_violations:
                    content_violation_parts.append(f"'{agent_name}' ({', '.join(content_violations)})")
                if extraction_failures:
                    extraction_failure_parts.append(f"'{agent_name}' ({', '.join(extraction_failures)})")

            messages = []

            # Emit event for content violations if any
            if content_violation_parts:
                msg = f"Cannot proceed: some agents have knowledge base documents blocked by content policy: {'; '.join(content_violation_parts)}"
                messages.append(msg)
                evt = {
                    "eventKind": EventKind.GUARDRAILS_VIOLATION,
                    "eventData": {
                        "blocked_agents": blocked_agents,
                        "message": msg,
                    },
                }
                pcb(evt)

            # Emit event for extraction failures if any
            if extraction_failure_parts:
                msg = f"Cannot proceed: text extraction failed for knowledge base documents in some agents: {'; '.join(extraction_failure_parts)}"
                messages.append(msg)
                evt = {
                    "eventKind": EventKind.GUARDRAILS_EXTRACTION_FAILED,
                    "eventData": {
                        "blocked_agents": blocked_agents,
                        "message": msg,
                    },
                }
                pcb(evt)

            combined_msg = " ".join(messages)
            if tcb:
                tcb(combined_msg)
            return combined_msg

    def _check_agent_runtime_guardrails(
        self,
        agent_id: str,
        documents: list,
        use_published: bool,
        pcb,
    ) -> dict:
        """
        Check agent documents against guardrails at runtime.

        This re-validates documents in case the guardrails pattern has changed
        since the documents were indexed.

        Args:
            agent_id: The agent ID
            documents: List of document dicts
            use_published: Whether to use published zone
            pcb: Push callback for events

        Returns:
            Dict with 'content_violations', 'extraction_failures', 'passed', and 'has_violations' keys
        """
        from backend.services.guardrails_service import check_agent_guardrails_at_runtime

        try:
            # Emit event that we're checking guardrails
            active_docs = [
                d
                for d in documents
                if isinstance(d, dict)
                and d.get("source") not in {"connection_folder"}
                and d.get("active")
                and not d.get("deletePending")
            ]
            if active_docs:
                pcb(
                    {
                        "eventKind": EventKind.GUARDRAILS_CHECKS,
                        "eventData": {
                            "message": f"Checking {len(active_docs)} knowledge base document(s) against content policy.",
                        },
                    },
                    store_event=True,
                )

            result = check_agent_guardrails_at_runtime(agent_id, documents, use_published)

            if result.get("has_violations"):
                content_violations = len(result.get("content_violations", []))
                extraction_failures = len(result.get("extraction_failures", []))
                total_blocked = content_violations + extraction_failures
                logger.warning(
                    f"Runtime guardrails check failed for agent {agent_id}: "
                    f"{total_blocked} documents blocked ({content_violations} content violations, {extraction_failures} extraction failures)"
                )

            return result
        except Exception as e:
            logger.exception(f"Error during runtime guardrails check for agent {agent_id}: {e}")
            # On error, allow the agent to proceed (fail open)
            return {"content_violations": [], "extraction_failures": [], "passed": documents, "has_violations": False}
