from __future__ import annotations

import asyncio
import json
import logging
import time
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional

from app.services.talkback_llm import (
    BROAD_SCOPE_CLARIFICATION_MESSAGE,
    BROAD_SCOPE_PROCESSING_CONSENT_TYPE,
    DEFAULT_FILE_RECONCILE_CLARIFICATION_MESSAGE,
    PRIOR_DF_FOLLOWUP_FILE_RECONCILE,
    _talkback_classify_context_relevance,
    _talkback_generate_contextual_question_smart,
    _talkback_is_broad_or_underconstrained_query,
    _talkback_merge_captured_question_time_fragment,
    _talkback_suggest_intent_and_agents,
)

from app.services.talkback.types import (
    ExecuteTalkback,
    GetAskmeService,
    NodeStatus,
    ResolveAuthorizedAgentIds,
    RunInThreadpool,
    TalkbackState,
)
from app.services.talkback.utils import (
    _clone_ask_request,
    _coerce_route_handler_result_to_dict,
    _ensure_web_prefix_for_search_web_request,
    _looks_like_clarification_reply,
    _safe_str,
)
from app.services.talkback.meta import (
    _agent_display_labels,
    _enriched_answer_for_context,
    _rewrite_context_should_omit_assistant_answer,
    _turn_for_llm_context,
    build_talkback_answer_meta,
)
from app.services.talkback.routing import (
    OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE,
    OUT_OF_CONTEXT_OPTION_LABEL_KEEP_CONTEXT,
    OUT_OF_CONTEXT_OPTION_LABEL_START_FRESH,
    _agent_scope_list_from_session,
    _build_intent_routing_question,
    _derive_request_id,
    _derive_talkback_db_session_key,
    _filter_agent_ids_by_datasource,
    _history_last_turn_expected_agent_disambiguation_reply,
    _infer_talkback_datasource_from_dispatch_and_intent,
    _last_turn_is_clarification_response_matching_question,
    _merge_original_question_with_clarification_selection,
    _out_of_context_clear_warrants_agent_handoff,
    _print_node,
    _reply_means_out_of_context_proceed,
    _reply_means_use_current_context_agents,
    _resolve_clarification_via_persisted_agent_scope,
    _resolve_talkback_file_agent_id,
)

logger = logging.getLogger(__name__)

OUT_OF_CONTEXT_CLEAR_CONFIRM_MESSAGE = (
    "This question looks unrelated to your recent conversation. "
    "Please select one option to continue."
)

OUT_OF_CONTEXT_CLARIFICATION_TYPE = "out_of_context_thread_selection"


def _out_of_context_clear_confirm_details() -> Dict[str, Any]:
    """``details`` shape aligned with file-upload clarification (``needs_clarification`` + options)."""
    return {
        "needs_clarification": True,
        "consent_type": OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE,
        "clarification_type": OUT_OF_CONTEXT_CLARIFICATION_TYPE,
        "clarification_message": OUT_OF_CONTEXT_CLEAR_CONFIRM_MESSAGE,
        "clarification_options": [
            {
                "id": "start_fresh",
                "label": OUT_OF_CONTEXT_OPTION_LABEL_START_FRESH,
                "action": "start_fresh",
            },
            {
                "id": "keep_context",
                "label": OUT_OF_CONTEXT_OPTION_LABEL_KEEP_CONTEXT,
                "action": "keep_context",
            },
        ],
    }


_UQL_CACHE_NODE_NAME = "UserQuestionLogCacheHitNode"
_PERSIST_TURN_NODE_NAME = "PersistTurnNode"


def _question_requires_file_only_dispatch(state: TalkbackState) -> bool:
    """
    True when the user question uses a document prefix (@file / @find / @analyze / @analyse).

    Uses ``state.file_cmd_parse`` when set (mixed path); otherwise parses the current question text
    so behavior stays correct even when ``FileCommandDetectNode`` returned early without populating
    ``file_cmd_parse`` on ``state``.
    """
    p = state.file_cmd_parse
    if isinstance(p, dict) and p.get("has_file_cmd"):
        return True
    try:
        from app.routes.askme import _detect_file_commands

        q = (
            (state.question_text or "").strip()
            or (state.talkback_raw_question_text or "").strip()
            or (str(getattr(state.request, "question", None) or "").strip())
        )
        if q and _detect_file_commands(q).get("has_file_cmd"):
            return True
    except Exception:
        pass
    return False


def _restrict_dispatch_agent_ids_for_file_datasource(
    state: TalkbackState,
    dispatch_agent_ids: List[str],
    scope_ids: List[str],
) -> List[str]:
    """
    For document/file commands, ``dispatch_agent_ids`` must only list agents with FILE datasource.

    When there is also a non-file clause (mixed intent), stash the pre-filter list on ``state``
    so the DB/API branch can still run against the router's agents.
    """
    if not _question_requires_file_only_dispatch(state):
        return list(dispatch_agent_ids)

    parse = state.file_cmd_parse if isinstance(state.file_cmd_parse, dict) else {}
    prior = [str(x).strip() for x in (dispatch_agent_ids or []) if str(x).strip()]
    pool: List[str] = list(prior)
    if not pool:
        pool = [str(x).strip() for x in (scope_ids or []) if str(x).strip()]
    if not pool:
        pool = [str(x).strip() for x in (state.resolved_agent_ids or []) if str(x).strip()]

    file_only = _filter_agent_ids_by_datasource(pool, "File")
    if not file_only:
        file_only = _filter_agent_ids_by_datasource(list(state.resolved_agent_ids or []), "File")

    if parse.get("has_non_file_clause") and prior:
        try:
            setattr(state, "_non_file_dispatch_agent_ids_for_mixed", prior)
        except Exception:
            pass

    # Never fall back to non-file agents here — that re-introduces DB/API ids into dispatch.
    return list(file_only) if file_only else []


class BaseNode(ABC):
    @property
    @abstractmethod
    def name(self) -> str:
        raise NotImplementedError

    @abstractmethod
    async def process(self, state: TalkbackState) -> TalkbackState:
        raise NotImplementedError


class _TalkbackStop(Exception):
    """Internal control flow for early returns (e.g., clarification response)."""


class _UserQuestionLogCacheHitNode(BaseNode):
    """Before routing: reuse a prior successful ``user_question_log`` row for the same normalized question."""

    def __init__(
        self,
        resolve_authorized_agent_ids: ResolveAuthorizedAgentIds,
        run_in_threadpool: RunInThreadpool,
        execute: ExecuteTalkback,
    ):
        self._resolve = resolve_authorized_agent_ids
        self._run_in_threadpool = run_in_threadpool
        self._execute = execute

    @property
    def name(self) -> str:
        return _UQL_CACHE_NODE_NAME

    async def process(self, state: TalkbackState) -> TalkbackState:
        req = state.request
        if getattr(req, "user_question_log_cache_injected", False):
            return state
        if bool(getattr(req, "is_dashboard", False)):
            return state
        if bool(getattr(req, "is_file", False)):
            return state
        if getattr(req, "question_id", None):
            return state

        raw_q = (getattr(req, "question", None) or "").strip()
        if not raw_q and getattr(req, "my_question", None):
            raw_q = str(getattr(req, "my_question")).strip()
        if not raw_q:
            return state

        from app.routes.askme import (
            _detect_file_commands,
            _has_meaningful_request_api,
            _has_meaningful_saved_sql,
            _infer_ask_datasource,
            _normalize_datasource_str,
        )
        from app.services.user_question_log_cache import (
            canonical_agent_id_for_cache,
            coerce_api_field_for_request,
            coerce_sql_field_for_request,
            find_reusable_log_row,
            normalize_question_for_cache,
            parse_agent_ids_from_log_column,
        )

        if _has_meaningful_saved_sql(getattr(req, "sql", None)):
            return state
        if _has_meaningful_request_api(getattr(req, "api_endpoint", None), getattr(req, "api_params", None)):
            return state
        try:
            if _detect_file_commands(raw_q).get("has_file_cmd"):
                return state
        except Exception:
            pass
        ds_raw = getattr(req, "datasource", None)
        if isinstance(ds_raw, str) and ds_raw.strip().lower() in ("file", "document"):
            return state

        qn = normalize_question_for_cache(raw_q)
        if not qn:
            return state

        def _lookup():
            return find_reusable_log_row(question_normalized=qn, user_id=state.current_user_id)

        row = await self._run_in_threadpool(_lookup)
        if row is None:
            return state

        row_agents = parse_agent_ids_from_log_column(row.agent_id)
        if not row_agents:
            return state

        client_agents_raw = getattr(req, "agent_id", None)
        if isinstance(client_agents_raw, list):
            has_client = bool([x for x in client_agents_raw if str(x).strip()])
        else:
            has_client = bool(str(client_agents_raw or "").strip())
        if has_client:
            if canonical_agent_id_for_cache(client_agents_raw) != canonical_agent_id_for_cache(row.agent_id):
                return state

        authorized_all = self._resolve(None, state.current_user_id)
        auth_set = set(authorized_all)
        use_agents = [a for a in row_agents if a in auth_set]
        if not use_agents:
            return state

        sql_req = coerce_sql_field_for_request(row.sql)
        ep = coerce_api_field_for_request(row.api_endpoint)
        pr = coerce_api_field_for_request(row.api_params)
        val = coerce_api_field_for_request(row.api_response_validator)
        sq_cached = None
        try:
            raw_sq = getattr(row, "sql_query_summaries", None)
            if isinstance(raw_sq, str) and raw_sq.strip():
                sq_cached = json.loads(raw_sq)
            elif raw_sq is not None:
                sq_cached = raw_sq
        except Exception:
            sq_cached = None

        updates: Dict[str, Any] = {
            "question": raw_q,
            "agent_id": list(use_agents) if len(use_agents) > 1 else use_agents[0],
            "user_question_log_cache_injected": True,
        }
        if sql_req is not None:
            updates["sql"] = sql_req
        if ep is not None:
            updates["api_endpoint"] = ep
        if pr is not None:
            updates["api_params"] = pr
        if val is not None:
            updates["api_response_validator"] = val
        if sq_cached:
            updates["sql_query_summaries_cached"] = sq_cached

        inferred = _infer_ask_datasource(
            _normalize_datasource_str(getattr(req, "datasource", None)),
            ep if ep is not None else row.api_endpoint,
            pr if pr is not None else row.api_params,
            multi_agent_unset_ok=True,
        )
        if inferred:
            updates["datasource"] = inferred

        cloned = _clone_ask_request(req, updates)
        try:
            setattr(cloned, "user_id", state.current_user_id)
        except Exception:
            pass

        raw = await self._execute(cloned, state.current_user)
        state.result = _coerce_route_handler_result_to_dict(raw)
        state.result["intent_routing"] = {
            "intent": "user_question_log_cache",
            "user_question_log_cache_hit": True,
        }
        state.result["context_cleared"] = state.context_cleared
        if state.context_cleared:
            state.result["context_clear_reason"] = state.context_clear_reason
        state.result["trace_id"] = state.trace.trace_id

        state.question_text = raw_q
        state.talkback_raw_question_text = raw_q
        if state.question_text:
            setattr(state.request, "question", state.question_text)
        setattr(state.request, "inContext", True)
        state.resolved_agent_ids = list(use_agents)
        state.dispatch_agent_ids = list(use_agents)
        state.intent_routing = dict(state.result["intent_routing"])
        state._uql_cache_hit_short_circuit = True
        return state


class _ResolveAgentsNode(BaseNode):
    def __init__(self, resolver: ResolveAuthorizedAgentIds):
        self._resolver = resolver

    @property
    def name(self) -> str:
        return "ResolveAgentsNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        raw_q = (getattr(state.request, "question", None) or "").strip()
        if not raw_q and getattr(state.request, "my_question", None):
            raw_q = str(getattr(state.request, "my_question")).strip()
        state.question_text = raw_q if raw_q else ""
        state.talkback_raw_question_text = raw_q if raw_q else ""

        # Ensure downstream execution steps (single/multi handlers) always see `request.question`.
        if state.question_text:
            setattr(state.request, "question", state.question_text)

        state.resolved_agent_ids = self._resolver(getattr(state.request, "agent_id", None), state.current_user_id)
        if not state.resolved_agent_ids:
            state.response = {
                "success": False,
                "message": "No authorized agents found for this user or requested agent(s)",
                "details": {},
                "trace_id": state.trace.trace_id,
            }
        return state


class _NormalizeDatasourceNode(BaseNode):
    @property
    def name(self) -> str:
        return "NormalizeDatasourceNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        # Talkback always operates in context mode.
        setattr(state.request, "inContext", True)

        ds = getattr(state.request, "datasource", None)
        state.datasource = ds
        if isinstance(ds, str):
            low = ds.strip().lower()
            if low in {"db", "database"}:
                setattr(state.request, "datasource", "Database")
                state.datasource = "Database"
            elif low == "api":
                setattr(state.request, "datasource", "API")
                state.datasource = "API"
            elif low in {"file", "document"}:
                setattr(state.request, "datasource", "File")
                state.datasource = "File"
        return state


class _FileCommandDetectNode(BaseNode):
    def __init__(self, run_in_threadpool: RunInThreadpool):
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "FileCommandDetectNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        question_text = (state.question_text or "").strip()
        if not question_text:
            return state

        # Local imports avoid circular import at module import-time.
        from app.routes.askme import _detect_file_commands, _ask_file_core

        parse = _detect_file_commands(question_text)
        if not parse.get("has_file_cmd"):
            return state

        if parse.get("has_conflict"):
            state.response = {
                "success": False,
                "message": "Conflicting command keywords are not allowed in the same query.",
                "details": {"error_code": "conflict_keywords", "question": question_text},
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()

        explicit_ds = str(state.datasource or "").strip().upper()
        if explicit_ds in ("DATABASE", "API"):
            state.response = {
                "success": False,
                "message": (
                    f"Datasource override '{state.datasource}' conflicts with file command in question. "
                    "Remove datasource override or split into separate queries."
                ),
                "details": {"error_code": "datasource_conflict", "question": question_text},
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()

        file_capable_ids = _filter_agent_ids_by_datasource(state.resolved_agent_ids, "File")
        if not file_capable_ids:
            picked = _resolve_talkback_file_agent_id(
                state, list(state.resolved_agent_ids or [])
            )
            if isinstance(picked, list):
                file_capable_ids = picked
            elif picked:
                file_capable_ids = [picked]
            else:
                file_capable_ids = []

        if not file_capable_ids:
            state.response = {
                "success": False,
                "message": "No file-capable agents are available for this request.",
                "details": {
                    "needs_clarification": True,
                    "clarification_message": "Please select an agent that supports file/document datasource.",
                    "error_code": "file_agent_required",
                },
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()

        if not parse.get("has_non_file_clause"):
            state.file_cmd_parse = parse
            state.dispatch_agent_ids = [str(x) for x in file_capable_ids if str(x).strip()]
            # Multi-agent support: pass the list if multiple, otherwise the single ID
            target_agent_id = file_capable_ids if len(file_capable_ids) > 1 else file_capable_ids[0]
            file_result = await _ask_file_core(
                user_id=state.current_user_id,
                agent_id=target_agent_id,
                question=parse.get("file_clause") or question_text,
                current_user=state.current_user,
                skip_log=False,
            )
            result = _coerce_route_handler_result_to_dict(file_result)
            result["intent_routing"] = state.intent_routing
            result["context_cleared"] = state.context_cleared
            if state.context_cleared:
                result["context_clear_reason"] = state.context_clear_reason
            result["trace_id"] = state.trace.trace_id
            if isinstance(result, dict):
                result["dispatch_agent_ids"] = list(state.dispatch_agent_ids)
            state.response = result

            # Refine dispatch scope from search results before stopping
            _update_dispatch_from_result(state)
            
            raise _TalkbackStop()

        # Mixed intent: keep file parse and let normal pipeline route non-file part (DB/API).
        state.file_cmd_parse = parse
        non_file_clause = str(parse.get("non_file_clause") or "").strip()
        if non_file_clause:
            setattr(state.request, "question", non_file_clause)
            state.question_text = non_file_clause
        state.datasource = None
        setattr(state.request, "datasource", None)
        return state


class _LoadSessionNode(BaseNode):
    def __init__(self, get_service: GetAskmeService, run_in_threadpool: RunInThreadpool):
        self._get_service = get_service
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "LoadSessionNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        if not state.question_text or not state.resolved_agent_ids:
            return state

        service_key = state.resolved_agent_ids if len(state.resolved_agent_ids) > 1 else state.resolved_agent_ids[0]
        service = self._get_service(agent_id=service_key)
        state.service = service
        if not service:
            return state

        session_key = _derive_talkback_db_session_key(service, state)
        state.session_key = session_key
        session_state = await self._run_in_threadpool(service.get_talkback_session_state, session_key)
        state.session_state = session_state

        # Load the previous agent_scope so DispatchDecisionNode can auto-select
        # the common agent when new routing overlaps with the prior turn's scope.
        state.previous_agent_scope = _agent_scope_list_from_session(session_state)

        history_turns: List[Dict[str, Any]] = []
        if session_state and session_state.get("history_json"):
            try:
                parsed = json.loads(session_state.get("history_json") or "[]")
                if isinstance(parsed, list):
                    history_turns = parsed
            except Exception:
                history_turns = []

        consent_ctx: Dict[str, Any] = {}
        if session_state and session_state.get("clarification_context_json"):
            try:
                cc = json.loads(session_state.get("clarification_context_json") or "{}")
                if isinstance(cc, dict):
                    consent_ctx = cc
            except Exception:
                consent_ctx = {}

        pending = bool(session_state and session_state.get("pending_clarification", 0))

        # Legacy: prior deployments used ``time_period_ambiguity``; clear stale pending so talkback does not mis-route.
        if (
            pending
            and str(consent_ctx.get("consent_type") or "").strip() == "time_period_ambiguity"
            and service
            and session_key
        ):
            await self._run_in_threadpool(service.clear_talkback_pending_clarification_only, session_key)
            session_state = await self._run_in_threadpool(service.get_talkback_session_state, session_key)
            state.session_state = session_state
            consent_ctx = {}
            if session_state and session_state.get("clarification_context_json"):
                try:
                    cc = json.loads(session_state.get("clarification_context_json") or "{}")
                    if isinstance(cc, dict):
                        consent_ctx = cc
                except Exception:
                    consent_ctx = {}
            pending = bool(session_state and session_state.get("pending_clarification", 0))

        # --- Broad-scope processing hint (session-stored; follow-up: "proceed" or merge with captured question) ---
        if (
            pending
            and consent_ctx.get("consent_type") == BROAD_SCOPE_PROCESSING_CONSENT_TYPE
            and service
            and session_key
        ):
            raw_q = (state.question_text or "").strip()
            captured = str(consent_ctx.get("captured_question") or "").strip()

            async def _reload_session() -> None:
                nonlocal session_state, consent_ctx, pending
                session_state = await self._run_in_threadpool(service.get_talkback_session_state, session_key)
                state.session_state = session_state
                consent_ctx = {}
                if session_state and session_state.get("clarification_context_json"):
                    try:
                        cc = json.loads(session_state.get("clarification_context_json") or "{}")
                        if isinstance(cc, dict):
                            consent_ctx = cc
                    except Exception:
                        consent_ctx = {}
                pending = bool(session_state and session_state.get("pending_clarification", 0))

            if raw_q:
                if raw_q.strip().lower() == "proceed":
                    restored = captured or raw_q
                    state.question_text = restored
                    setattr(state.request, "question", restored)
                    try:
                        setattr(state.request, "my_question", restored)
                    except Exception:
                        pass
                    await self._run_in_threadpool(service.clear_talkback_pending_clarification_only, session_key)
                    await _reload_session()
                    state.skip_broad_scope_check_once = True
                elif captured:
                    # Any other reply elaborates the stored ask; merge and continue to routing / agent clarification.
                    def _merge_broad_followup() -> str:
                        return _talkback_merge_captured_question_time_fragment(
                            service,
                            captured,
                            raw_q,
                            temperature=0.0,
                            max_tokens=400,
                            use_session=False,
                        )

                    merged = (await self._run_in_threadpool(_merge_broad_followup)).strip()
                    if not merged:
                        merged = f"{captured} {raw_q}".strip()
                    state.question_text = merged
                    setattr(state.request, "question", merged)
                    try:
                        setattr(state.request, "my_question", merged)
                    except Exception:
                        pass
                    await self._run_in_threadpool(service.clear_talkback_pending_clarification_only, session_key)
                    await _reload_session()
                    state.skip_broad_scope_check_once = True
                else:
                    state.response = {
                        "success": True,
                        "message": "Scope clarification",
                        "details": {
                            "needs_clarification": True,
                            "clarification_message": BROAD_SCOPE_CLARIFICATION_MESSAGE,
                            "consent_type": BROAD_SCOPE_PROCESSING_CONSENT_TYPE,
                            "clarification_agent_ids": [],
                        },
                        "intent_routing": state.intent_routing if isinstance(state.intent_routing, dict) else {},
                        "context_cleared": state.context_cleared,
                        "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                        "trace_id": state.trace.trace_id,
                    }
                    raise _TalkbackStop()

        # --- Out-of-context session clear confirmation (session-stored; follow-up: proceed vs use current) ---
        if (
            pending
            and consent_ctx.get("consent_type") == OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE
            and service
            and session_key
        ):
            raw_q = (state.question_text or "").strip()
            captured = str(consent_ctx.get("captured_question") or "").strip()

            async def _reload_session_ooc() -> None:
                nonlocal session_state, consent_ctx, pending, history_turns
                session_state = await self._run_in_threadpool(service.get_talkback_session_state, session_key)
                state.session_state = session_state
                consent_ctx = {}
                if session_state and session_state.get("clarification_context_json"):
                    try:
                        cc = json.loads(session_state.get("clarification_context_json") or "{}")
                        if isinstance(cc, dict):
                            consent_ctx = cc
                    except Exception:
                        consent_ctx = {}
                pending = bool(session_state and session_state.get("pending_clarification", 0))
                history_turns = []
                if session_state and session_state.get("history_json"):
                    try:
                        parsed = json.loads(session_state.get("history_json") or "[]")
                        if isinstance(parsed, list):
                            history_turns = parsed
                    except Exception:
                        history_turns = []

            prior_snap_raw = consent_ctx.get("prior_agent_scope")
            prior_snap: List[str] = (
                [str(x).strip() for x in prior_snap_raw if str(x).strip()]
                if isinstance(prior_snap_raw, list)
                else []
            )
            clar_reason = str(consent_ctx.get("context_clear_reason") or "").strip()

            if raw_q and _reply_means_out_of_context_proceed(raw_q):
                await self._run_in_threadpool(
                    service.clear_talkback_session_state,
                    session_key,
                    list(state.resolved_agent_ids),
                )
                await _reload_session_ooc()
                state.history_turns = history_turns
                state.previous_agent_scope = _agent_scope_list_from_session(session_state)
                state.context_cleared = True
                state.context_clear_reason = clar_reason or "User confirmed starting a fresh thread."
                if prior_snap:
                    state.pre_clear_dispatch_scope = list(prior_snap)
                restored = captured or raw_q
                state.question_text = restored
                setattr(state.request, "question", restored)
                try:
                    setattr(state.request, "my_question", restored)
                except Exception:
                    pass
                state.routing_agent_scope_ids = None
                state.pending_prior_df_followup = None
            elif raw_q and _reply_means_use_current_context_agents(raw_q):
                await self._run_in_threadpool(service.clear_talkback_pending_clarification_only, session_key)
                await _reload_session_ooc()
                state.history_turns = history_turns
                state.previous_agent_scope = _agent_scope_list_from_session(session_state)
                state.context_cleared = False
                state.context_clear_reason = ""
                restored = captured or raw_q
                state.question_text = restored
                setattr(state.request, "question", restored)
                try:
                    setattr(state.request, "my_question", restored)
                except Exception:
                    pass
                state.skip_context_relevance_once = True
                resolved_set = {str(x) for x in state.resolved_agent_ids}
                narrowed = [str(x) for x in prior_snap if str(x) in resolved_set]
                state.routing_agent_scope_ids = narrowed if narrowed else (prior_snap or None)
            elif not raw_q:
                state.response = {
                    "success": True,
                    "message": "Clarification needed",
                    "details": _out_of_context_clear_confirm_details(),
                    "intent_routing": state.intent_routing if isinstance(state.intent_routing, dict) else {},
                    "context_cleared": state.context_cleared,
                    "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                    "trace_id": state.trace.trace_id,
                }
                raise _TalkbackStop()
            elif _looks_like_clarification_reply(raw_q) and not (
                _reply_means_out_of_context_proceed(raw_q)
                or _reply_means_use_current_context_agents(raw_q)
            ):
                state.response = {
                    "success": True,
                    "message": "Clarification needed",
                    "details": _out_of_context_clear_confirm_details(),
                    "intent_routing": state.intent_routing if isinstance(state.intent_routing, dict) else {},
                    "context_cleared": state.context_cleared,
                    "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                    "trace_id": state.trace.trace_id,
                }
                raise _TalkbackStop()

        llm_df_pending = pending and consent_ctx.get("consent_type") == "llm_prior_df"
        # Broad-scope hint uses the same pending flag but is not agent disambiguation.
        disambig_pending = pending and consent_ctx.get("consent_type") not in (
            "llm_prior_df",
            BROAD_SCOPE_PROCESSING_CONSENT_TYPE,
            "time_period_ambiguity",
            OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE,
        )
        search_web = bool(getattr(state.request, "search_web", False))
        looks_like_clar_reply = _looks_like_clarification_reply(state.question_text)
        # Preserve mixed DB/API+FILE decomposition across clarification turns.
        # Clarification replies are usually short ("agent 1", "both") and do not include
        # file commands, so we restore prior file parse from session context.
        if pending and looks_like_clar_reply:
            saved_parse = consent_ctx.get("file_cmd_parse")
            if isinstance(saved_parse, dict) and saved_parse.get("has_file_cmd"):
                state.file_cmd_parse = saved_parse

        # Narrow intent-routing scope ONLY when we're genuinely continuing a pending disambiguation.
        if disambig_pending and looks_like_clar_reply:
            raw_cids = consent_ctx.get("clarification_agent_ids")
            if isinstance(raw_cids, list) and raw_cids:
                resolved_set = {str(x) for x in state.resolved_agent_ids}
                narrowed = [str(x) for x in raw_cids if str(x) in resolved_set]
                if narrowed:
                    state.routing_agent_scope_ids = narrowed

        # If last turn asked for clarification, treat current input as a clarification response and persist it.
        if pending:
            if llm_df_pending and search_web:
                pf = str(consent_ctx.get("prior_df_followup") or "").strip().lower()
                if pf in ("askme_analysis", "web_enrichment"):
                    state.pending_prior_df_followup = pf
                await self._run_in_threadpool(service.clear_talkback_pending_clarification_only, session_key)
                session_state = await self._run_in_threadpool(service.get_talkback_session_state, session_key)
                state.session_state = session_state
            elif not llm_df_pending:
                if not looks_like_clar_reply:
                    # User asked a fresh question while a prior disambiguation was pending.
                    # Drop pending clarification so routing can use full resolved agent scope.
                    await self._run_in_threadpool(service.clear_talkback_pending_clarification_only, session_key)
                    session_state = await self._run_in_threadpool(service.get_talkback_session_state, session_key)
                    state.session_state = session_state
                    state.routing_agent_scope_ids = None
                    consent_ctx = {}
                else:
                    if consent_ctx.get("consent_type") == BROAD_SCOPE_PROCESSING_CONSENT_TYPE:
                        state.response = {
                            "success": True,
                            "message": "Scope clarification",
                            "details": {
                                "needs_clarification": True,
                                "clarification_message": BROAD_SCOPE_CLARIFICATION_MESSAGE,
                                "consent_type": BROAD_SCOPE_PROCESSING_CONSENT_TYPE,
                                "clarification_agent_ids": [],
                            },
                            "intent_routing": state.intent_routing if isinstance(state.intent_routing, dict) else {},
                            "context_cleared": state.context_cleared,
                            "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                            "trace_id": state.trace.trace_id,
                        }
                        raise _TalkbackStop()
                    if consent_ctx.get("consent_type") == OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE:
                        state.response = {
                            "success": True,
                            "message": "Clarification needed",
                            "details": _out_of_context_clear_confirm_details(),
                            "intent_routing": state.intent_routing if isinstance(state.intent_routing, dict) else {},
                            "context_cleared": state.context_cleared,
                            "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                            "trace_id": state.trace.trace_id,
                        }
                        raise _TalkbackStop()
                    state.agent_disambiguation_follow_up = True
                    request_id = _derive_request_id(state)
                    history_turns.append(
                        {
                            "question": f"Clarification response: {state.question_text}",
                            "answer": "",
                            "ts": time.strftime("%Y-%m-%dT%H:%M:%S"),
                            "request_id": request_id,
                        }
                    )
                    scope_ids = _agent_scope_list_from_session(session_state) or list(state.resolved_agent_ids)
                    preserved_cc: Dict[str, Any] = {}
                    if consent_ctx.get("consent_type") != "llm_prior_df":
                        pids = consent_ctx.get("clarification_agent_ids")
                        if isinstance(pids, list) and pids:
                            preserved_cc["clarification_agent_ids"] = [str(x) for x in pids if str(x).strip()]
                            preserved_cc["consent_type"] = "agent_disambiguation"
                    await self._run_in_threadpool(
                        service.upsert_talkback_session_state,
                        session_key,
                        state.current_user_id,
                        state.resolved_agent_ids,
                        history_turns[-10:],
                        session_state.get("openai_response_id") if session_state else None,
                        False,
                        None,
                        preserved_cc if preserved_cc else {},
                        scope_ids,
                    )
                    refreshed_ss = await self._run_in_threadpool(
                        service.get_talkback_session_state, session_key
                    )
                    if refreshed_ss:
                        state.session_state = refreshed_ss
                        session_state = refreshed_ss

        state.history_turns = history_turns

        # If DB has no agent_scope yet (missing row, or "[]" from legacy clears), seed from the request's
        # authorized agents so clients and DispatchDecisionNode see a non-blank saved scope.
        latest_for_scope = state.session_state
        scope_now = _agent_scope_list_from_session(latest_for_scope)
        if (not scope_now) and state.resolved_agent_ids:
            refreshed = await self._run_in_threadpool(
                service.upsert_talkback_session_state,
                session_key,
                state.current_user_id,
                state.resolved_agent_ids,
                None,
                None,
                None,
                None,
                None,
                None,
            )
            state.session_state = refreshed
            state.previous_agent_scope = _agent_scope_list_from_session(refreshed)

        # Session context for routing/execution
        if session_state and session_state.get("openai_response_id"):
            state.session_context = {"openai_response_id": session_state.get("openai_response_id")}
        else:
            state.session_context = {}

        return state


class _ContextRelevanceNode(BaseNode):
    def __init__(self, run_in_threadpool: RunInThreadpool):
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "ContextRelevanceNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        if not state.service or not state.session_key:
            return state

        if getattr(state, "skip_context_relevance_once", False):
            state.skip_context_relevance_once = False
            _print_node(
                "skip",
                self.name,
                state.trace.trace_id,
                state,
                reason="skip_context_relevance_once",
            )
            return state

        if not state.history_turns:
            return state

        if getattr(state, "agent_disambiguation_follow_up", False):
            _print_node(
                "skip",
                self.name,
                state.trace.trace_id,
                state,
                reason="agent_disambiguation_follow_up",
            )
            return state
        if _last_turn_is_clarification_response_matching_question(state):
            _print_node(
                "skip",
                self.name,
                state.trace.trace_id,
                state,
                reason="clarification_response_row_matches_question",
            )
            return state

        # Also skip if it's a direct response to the "exact vs similar" file search clarification
        if _is_file_search_clarification_response(state):
            _print_node(
                "skip",
                self.name,
                state.trace.trace_id,
                state,
                reason="file_search_clarification_response",
            )
            return state

        _print_node(
            "call",
            self.name,
            state.trace.trace_id,
            state,
            fn="classify_context_relevance",
            previous_turns=len(state.history_turns[-10:]),
            current_question=_safe_str(state.question_text),
        )
        prior_for_llm: List[Dict[str, Any]] = []
        for t in state.history_turns:
            if isinstance(t, dict):
                prior_for_llm.append(_turn_for_llm_context(t))
            else:
                prior_for_llm.append({"question": "", "answer": str(t)})
        relation = await self._run_in_threadpool(
            _talkback_classify_context_relevance,
            state.service,
            prior_for_llm,
            state.question_text,
        )
        _print_node(
            "result",
            self.name,
            state.trace.trace_id,
            state,
            is_related=relation.get("is_related") if isinstance(relation, dict) else None,
            confidence=relation.get("confidence") if isinstance(relation, dict) else None,
            reason=_safe_str(relation.get("reason"), 240) if isinstance(relation, dict) else None,
        )
        if state.history_turns and not bool(relation.get("is_related", True)):
            if not _out_of_context_clear_warrants_agent_handoff(
                state.history_turns,
                state.intent_routing if isinstance(state.intent_routing, dict) else None,
            ):
                _print_node(
                    "skip",
                    self.name,
                    state.trace.trace_id,
                    state,
                    reason="out_of_context_suppressed_same_or_overlapping_dispatch",
                )
                return state
            reason = str(
                relation.get("reason", "") or "Current question is not related to previous context."
            )
            prior_scope_copy: List[str] = (
                list(state.previous_agent_scope) if state.previous_agent_scope else []
            )
            q = (state.question_text or "").strip()
            request_id = _derive_request_id(state)
            session_key = state.session_key
            dispatch_scope = (
                list(state.dispatch_agent_ids) if state.dispatch_agent_ids else list(state.resolved_agent_ids)
            )
            clar_meta: Dict[str, Any] = {
                "turn_type": "clarification",
                "needs_clarification": True,
                "consent_type": OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE,
            }
            await self._run_in_threadpool(
                state.service.append_turn_to_talkback_session,
                session_key,
                state.current_user_id,
                state.resolved_agent_ids,
                q,
                OUT_OF_CONTEXT_CLEAR_CONFIRM_MESSAGE,
                10,
                request_id,
                clar_meta,
                dispatch_scope,
            )
            cc_ctx: Dict[str, Any] = {
                "consent_type": OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE,
                "captured_question": q,
                "prior_agent_scope": prior_scope_copy,
                "context_clear_reason": reason,
            }
            if isinstance(state.file_cmd_parse, dict):
                cc_ctx["file_cmd_parse"] = state.file_cmd_parse
            await self._run_in_threadpool(
                state.service.upsert_talkback_session_state,
                session_key,
                state.current_user_id,
                state.resolved_agent_ids,
                None,
                None,
                True,
                OUT_OF_CONTEXT_CLEAR_CONFIRM_MESSAGE,
                cc_ctx,
                dispatch_scope,
            )
            state.response = {
                "success": True,
                "message": "Clarification needed",
                "details": _out_of_context_clear_confirm_details(),
                "intent_routing": state.intent_routing if isinstance(state.intent_routing, dict) else {},
                "context_cleared": False,
                "context_clear_reason": None,
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()
        return state


class _PreparePriorQuestionsNode(BaseNode):
    @property
    def name(self) -> str:
        return "PreparePriorQuestionsNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        prior_questions: List[str] = []
        for t in state.history_turns:
            if not isinstance(t, dict):
                continue
            q = str(t.get("question", "")).strip()
            if not q:
                continue
            if q.lower().startswith("clarification response:"):
                q = q.split(":", 1)[1].strip() if ":" in q else ""
                if not q:
                    continue
            prior_questions.append(q)

        if prior_questions:
            setattr(state.request, "inContext", True)
            setattr(state.request, "previous_questions", prior_questions[-10:])
        elif state.context_cleared:
            setattr(state.request, "inContext", True)
            setattr(state.request, "previous_questions", [])

        # If context was cleared, do not pass old OpenAI response id
        if state.context_cleared:
            state.session_context = {}

        return state


class _ContextRewriteForRoutingNode(BaseNode):
    def __init__(self, run_in_threadpool: RunInThreadpool):
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "ContextRewriteForRoutingNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        routing_text = state.question_text

        clar_tail = (
            state.history_turns
            and isinstance(state.history_turns[-1], dict)
            and str(state.history_turns[-1].get("question", "")).strip().lower().startswith("clarification response:")
        )
        if getattr(state, "agent_disambiguation_follow_up", False) or clar_tail:
            _merge_original_question_with_clarification_selection(state)
            state.routing_question_text = state.question_text
            return state

        prev_questions = getattr(state.request, "previous_questions", None)
        if state.service and getattr(state.request, "inContext", False) and prev_questions:
            try:
                # Provide richer context than question-only: include recent Q+A turns when available.
                # This is not gating/heuristics; it just supplies more context to the prompt.
                enriched: List[str] = []
                for t in (state.history_turns or [])[-10:]:
                    if not isinstance(t, dict):
                        continue
                    q = str(t.get("question", "")).strip()
                    if not q:
                        continue
                    if q.lower().startswith("clarification response:"):
                        continue
                    if _rewrite_context_should_omit_assistant_answer(t):
                        enriched.append(f"Q: {q}")
                        continue
                    a = _enriched_answer_for_context(str(t.get("answer", "")), t.get("answer_meta"))
                    if a:
                        enriched.append(f"Q: {q}\nA: {a}")
                    else:
                        enriched.append(f"Q: {q}")
                if not enriched:
                    enriched = list(prev_questions)

                _print_node(
                    "call",
                    self.name,
                    state.trace.trace_id,
                    state,
                    fn="generate_contextual_question_smart",
                    previous_items=len(enriched),
                    current_question=_safe_str(state.question_text),
                    use_session=False,
                )
                routing_text = await self._run_in_threadpool(
                    _talkback_generate_contextual_question_smart,
                    state.service,
                    enriched,
                    state.question_text,
                    use_session=False,
                )
            except Exception:
                routing_text = state.question_text

        # Promote LLM rewrite to the active question when context is kept (is_related path).
        # Trust the rewrite prompt: it returns the question unchanged when not context-dependent,
        # so we apply the result without additional heuristics. The rewritten question must drive
        # downstream routing/execution + history persistence, otherwise short follow-ups like
        # "for 2025" reach SQL generation without the prior turn's metric/intent.
        rewritten = (routing_text or "").strip()
        original_q = (state.question_text or "").strip()

        # Preserve `@askme` / `@web` command markers if the rewrite stripped them. Downstream
        # `_PriorDfConsentNode` and `_prior_df_apply_post_heuristics` pin prior_df_followup
        # from these literal tags; losing them re-routes prior-df / prediction prompts back
        # through SQL generation and the user's `df` is never used.
        raw_for_tags = (
            getattr(state, "talkback_raw_question_text", "") or original_q or ""
        ).lower()
        rewritten_lower = rewritten.lower()
        if rewritten:
            if "@web" in raw_for_tags and "@web" not in rewritten_lower:
                rewritten = f"@web {rewritten}"
            elif "@askme" in raw_for_tags and "@askme" not in rewritten_lower:
                rewritten = f"@askme {rewritten}"
            routing_text = rewritten

        promote = bool(
            rewritten
            and rewritten != original_q
            and not state.context_cleared
        )
        if promote:
            state.question_text = rewritten
            try:
                setattr(state.request, "question", rewritten)
                # Always sync my_question too, even if it was previously None, so downstream
                # handlers that fall back to my_question see the rewritten question.
                setattr(state.request, "my_question", rewritten)
            except Exception:
                pass
            state.talkback_history_question_text = rewritten
            _print_node(
                "promote_rewrite",
                self.name,
                state.trace.trace_id,
                state,
                original=_safe_str(original_q),
                rewritten=_safe_str(rewritten),
            )

        state.routing_question_text = routing_text
        return state


class _IntentRoutingNode(BaseNode):
    def __init__(self, run_in_threadpool: RunInThreadpool):
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "IntentRoutingNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        if not state.service or not state.resolved_agent_ids or not state.question_text:
            state.intent_routing = None
            return state

        routing_ids = (
            list(state.routing_agent_scope_ids)
            if state.routing_agent_scope_ids
            else list(state.resolved_agent_ids)
        )
        filtered_ids = _filter_agent_ids_by_datasource(routing_ids, state.datasource)
        if str(state.datasource or "").strip() and not filtered_ids:
            state.intent_routing = {
                "intent": "ambiguous",
                "suggested_agent_ids": [],
                "needs_clarification": True,
                "clarification_message": (
                    f"No agents in current scope support datasource '{state.datasource}'."
                ),
                "clarification_agent_ids": [],
                "prior_df_followup": "none",
                "misrouted_web": False,
            }
            return state
        routing_ids = filtered_ids or routing_ids

        # If is_file=True, restrict the agent pool to file-capable agents only before calling the LLM router.
        is_file_request = False
        try:
            is_file_request = bool(getattr(state.request, "is_file", False))
        except Exception:
            pass

        if is_file_request:
            file_only_ids = _filter_agent_ids_by_datasource(routing_ids, "File")
            if file_only_ids:
                routing_ids = file_only_ids

        parse_fc = state.file_cmd_parse if isinstance(state.file_cmd_parse, dict) else {}
        if _question_requires_file_only_dispatch(state) and not parse_fc.get("has_non_file_clause"):
            file_only_routing = _filter_agent_ids_by_datasource(routing_ids, "File")
            if file_only_routing:
                routing_ids = file_only_routing

        def run_routing() -> Dict[str, Any]:
            rq = _build_intent_routing_question(state)
            _print_node(
                "call",
                self.name,
                state.trace.trace_id,
                state,
                fn="suggest_intent_and_agents",
                question=_safe_str(rq),
                agent_ids=routing_ids,
                temperature=0.2,
                max_tokens=800,
                use_session=True,
                session_context_keys=list((state.session_context or {}).keys()),
            )
            return _talkback_suggest_intent_and_agents(
                state.service,
                rq,
                agent_ids=routing_ids,
                temperature=0.2,
                max_tokens=800,
                use_session=True,
                session_context=state.session_context,
            )

        try:
            state.intent_routing = await self._run_in_threadpool(run_routing)
        except Exception as e:
            state.intent_routing = {"error": str(e)}
        return state


class _DispatchDecisionNode(BaseNode):
    def __init__(self, get_service: GetAskmeService, run_in_threadpool: RunInThreadpool):
        self._get_service = get_service
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "DispatchDecisionNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        ir = state.intent_routing
        scope_ids = (
            list(state.routing_agent_scope_ids)
            if state.routing_agent_scope_ids
            else list(state.resolved_agent_ids)
        )

        # Re-read agent_scope from DB: LoadSession ran before ContextRelevance may have
        # cleared session_state; prefer DB. Treat empty list as "no persisted narrow scope".
        persisted_scope_for_clar: Optional[List[str]] = state.previous_agent_scope
        if state.service and state.session_key:
            try:
                fresh = await self._run_in_threadpool(
                    state.service.get_talkback_session_state, state.session_key
                )
                if fresh:
                    parsed = _agent_scope_list_from_session(fresh)
                    if parsed is not None:
                        persisted_scope_for_clar = parsed if parsed else None
            except Exception:
                pass
        if state.context_cleared:
            persisted_scope_for_clar = None

        if not ir or not isinstance(ir, dict) or ir.get("error"):
            # Routing failed (timeout, parse error, etc.) — safe fallback to first
            # agent only; never fan out blindly to every resolved agent.
            if scope_ids:
                if _question_requires_file_only_dispatch(state):
                    fo = _filter_agent_ids_by_datasource(list(scope_ids), "File")
                    state.dispatch_agent_ids = [str(fo[0])] if fo else []
                else:
                    state.dispatch_agent_ids = [str(scope_ids[0])]
            else:
                state.dispatch_agent_ids = []
            return state

        dispatch_agent_ids = list(scope_ids)

        prior_df = str(ir.get("prior_df_followup", "none")).lower()
        if prior_df in ("askme_analysis", "web_enrichment", "file_reconcile"):
            if scope_ids:
                if _question_requires_file_only_dispatch(state):
                    fo = _filter_agent_ids_by_datasource(list(scope_ids), "File")
                    dispatch_agent_ids = [str(fo[0])] if fo else []
                else:
                    dispatch_agent_ids = [str(scope_ids[0])]
            state.intent_routing = {
                **ir,
                "suggested_agent_ids": [],
            }
        else:
            suggested = ir.get("suggested_agent_ids") or []
            needs_clar = bool(ir.get("needs_clarification"))
            resolved_set = set(str(a) for a in scope_ids)
            intent_str = str(ir.get("intent", "")).strip().lower()
            # Resolve user's response to "exact or similar?" clarification regardless of
            # what the LLM classifies the short reply as (e.g. "similar files" may not
            # route as intent="document"). Check last clarification message first.
            _raw_choice = getattr(state, "disambiguation_user_choice", None)
            _q_lower_clar = (
                _raw_choice or state.talkback_raw_question_text or state.question_text
            ).strip().lower()
            # Find the most recent clarification still awaiting a user choice. Skip
            # ``"Clarification response: ..."`` stub turns (LoadSessionNode appends one
            # when the user picks an option) so we can locate the clarification just
            # above it. Stop the moment we hit a real answer turn — reaching past a
            # resolved turn would mis-classify a brand-new question as a clarification
            # response and dispatch it to all agents.
            _last_clar_msg = ""
            for _t in reversed(state.history_turns or []):
                if not isinstance(_t, dict):
                    continue
                _tq = str(_t.get("question", "")).strip()
                if _tq.lower().startswith("clarification response:"):
                    continue
                _meta = _t.get("answer_meta")
                if isinstance(_meta, dict) and _meta.get("needs_clarification"):
                    _last_clar_msg = str(_t.get("answer") or "").strip().lower()
                break

            if _last_clar_msg and ("find the exact file" in _last_clar_msg or "find similar files" in _last_clar_msg or "select one option to continue" in _last_clar_msg):
                ir["needs_clarification"] = False
                # Recover original question from the turn that issued the prompt
                _original_q = state.question_text
                for _t in reversed(state.history_turns):
                    if not isinstance(_t, dict):
                        continue
                    _meta = _t.get("answer_meta")
                    if isinstance(_meta, dict) and _meta.get("needs_clarification"):
                        _prev_ans = str(_t.get("answer") or "").strip().lower()
                        # Recover from the options step message
                        if "select one option to continue" in _prev_ans:
                            _oq = str(_t.get("question") or "").strip()
                            if _oq and not _oq.lower().startswith("clarification"):
                                _original_q = _oq
                            break

                if "exact" in _q_lower_clar or "@find" in _q_lower_clar or "first" in _q_lower_clar:
                    _prefix = "@find"
                elif "similar" in _q_lower_clar or "@file" in _q_lower_clar or "second" in _q_lower_clar:
                    _prefix = "@file"
                else:
                    _prefix = None

                if _prefix:
                    _prefixed = f"{_prefix} {_original_q}"
                    state.routing_question_text = _prefixed
                    state.question_text = _prefixed
                    state.talkback_history_question_text = _prefixed
                    try:
                        setattr(state.request, "question", _prefixed)
                        setattr(state.request, "my_question", _prefixed)
                    except Exception:
                        pass
                    try:
                        from app.routes.askme import _detect_file_commands
                        _parse = _detect_file_commands(_prefixed)
                        if _parse.get("has_file_cmd") and not _parse.get("has_conflict"):
                            state.file_cmd_parse = _parse
                            state.datasource = "File"
                            setattr(state.request, "datasource", "File")
                    except Exception:
                        pass

                state.intent_routing = ir
                needs_clar = False
                state.dispatch_agent_ids = _restrict_dispatch_agent_ids_for_file_datasource(
                    state, dispatch_agent_ids, scope_ids
                )
                return state

            # If explicit is_file flag is set, trigger the options clarification immediately
            # regardless of what the LLM classified the intent as.
            is_file_request = False
            try:
                is_file_request = bool(getattr(state.request, "is_file", False))
            except Exception:
                pass

            if is_file_request and not _is_file_search_clarification_response(state):
                ir["needs_clarification"] = True
                ir["clarification_message"] = "Choose how to proceed with the file search. Please select one option to continue."
                ir["is_custom_file_clarification"] = True
                ir["consent_type"] = "file_search_action"
                ir["clarification_type"] = "file_search_action_selection"
                ir["clarification_options"] = [
                    {
                        "id": "exact",
                        "label": "Find the exact file",
                        "action": "exact"
                    },
                    {
                        "id": "similar",
                        "label": "Find similar files",
                        "action": "similar"
                    }
                ]
                ir.pop("clarification_agent_ids", None)
                state.intent_routing = ir
                return state

            if needs_clar and scope_ids:
                # --- Resolve via persisted prior dispatch scope (talkback ``agent_scope``) ---
                # The subset rule ("if prev ⊆ clar, reuse prev; else ask") is the single source
                # of truth for both context-cleared and in-context follow-ups, and applies to
                # every intent (including ``multi_agent``) — needs_clarification=True always
                # means "user must pick", and a subset prev tells us they already did.
                clar_aids = ir.get("clarification_agent_ids") or []
                # Live previous_agent_scope is wiped by ContextRelevanceNode on context-clear;
                # pre_clear_dispatch_scope keeps the snapshot so identical / subset follow-ups
                # still reuse the prior agents in that case.
                prev_scope_for_resolve: Optional[List[str]] = (
                    list(state.previous_agent_scope)
                    if state.previous_agent_scope
                    else (
                        list(state.pre_clear_dispatch_scope)
                        if state.context_cleared and state.pre_clear_dispatch_scope
                        else None
                    )
                )
                auto_from_scope: Optional[List[str]] = (
                    _resolve_clarification_via_persisted_agent_scope(
                        prev_scope_for_resolve, clar_aids
                    )
                    if clar_aids
                    else None
                )
                if auto_from_scope:
                    dispatch_agent_ids = auto_from_scope
                    state.intent_routing = {
                        **ir,
                        "needs_clarification": False,
                        "suggested_agent_ids": auto_from_scope,
                        "clarification_agent_ids": auto_from_scope,
                        "clarification_message": None,
                    }
                    ir = state.intent_routing
                    logger.info(
                        "TalkbackOrchestrator reusing prior agent_scope %s "
                        "for clarification_agent_ids %s (intent=%s, context_cleared=%s, "
                        "prev ⊆ clar)",
                        auto_from_scope,
                        clar_aids,
                        intent_str or "<none>",
                        state.context_cleared,
                    )
                else:
                    # The subset rule already ran above; reaching here means prev_scope is
                    # NOT a subset of clar (some prev agent was dropped, prev was empty, or
                    # the router returned no clarification candidates). For context-cleared
                    # follow-ups we still surface a helpful merged-list message so the user
                    # can pick between continuing with the prior agent(s) or switching to a
                    # newly-suggested one. For all other paths we keep the LLM clarification
                    # as-is and dispatch a placeholder first agent.
                    clar_list = [str(a).strip() for a in clar_aids if str(a).strip()]
                    prev_scope_strs = list(prev_scope_for_resolve or [])
                    prev_scope_set = set(prev_scope_strs)
                    new_candidates = [
                        aid for aid in clar_list if aid not in prev_scope_set
                    ]
                    if (
                        state.context_cleared
                        and prev_scope_strs
                        and clar_list
                        and new_candidates
                    ):
                        merged = list(prev_scope_strs)
                        seen = set(merged)
                        for aid in new_candidates:
                            if aid in seen:
                                continue
                            merged.append(aid)
                            seen.add(aid)
                        prev_label = ", ".join(_agent_display_labels(prev_scope_strs))
                        new_label = ", ".join(_agent_display_labels(new_candidates))
                        state.intent_routing = {
                            **ir,
                            "needs_clarification": True,
                            "clarification_agent_ids": merged,
                            "clarification_message": (
                                "Your previous question was answered by "
                                f"{prev_label}. "
                                "For this question, you can continue with that agent or choose "
                                f"one of these: {new_label}. "
                                "Which agent should I use?"
                            ),
                        }
                        ir = state.intent_routing
                    # Clarification is still required; keep dispatch non-executable.
                    dispatch_agent_ids = []
            elif suggested:
                narrow = [str(a) for a in suggested if str(a) in resolved_set]
                if narrow:
                    dispatch_agent_ids = narrow

                    fan_out_all = (
                        intent_str != "multi_agent"
                        and len(scope_ids) > 1
                        and len(dispatch_agent_ids) == len(scope_ids)
                        and not needs_clar
                    )
                    non_multi_too_many = (
                        intent_str != "multi_agent"
                        and len(dispatch_agent_ids) > 1
                        and len(scope_ids) > 1
                    )

                    # --- Dynamic Intent/Agent Shift & Ambiguity Validation ---
                    # reference_scope: snapshot of what we were using before this question.
                    reference_scope = (
                        list(state.pre_clear_dispatch_scope)
                        if state.context_cleared and state.pre_clear_dispatch_scope
                        else (list(persisted_scope_for_clar) if persisted_scope_for_clar else None)
                    )

                    if reference_scope and intent_str != "multi_agent":
                        ref_set = set(str(x) for x in reference_scope)
                        common = [a for a in dispatch_agent_ids if a in ref_set]

                        # Subset/Overlap check: if exactly one agent from current suggestions
                        # matches the previous scope, we auto-select it to avoid annoying prompts.
                        if len(common) == 1 and len(dispatch_agent_ids) >= 1:
                            dispatch_agent_ids = common
                            logger.info(
                                "Talkback auto-selected agent %s from prior scope %s (context_cleared=%s)",
                                common[0], reference_scope, state.context_cleared
                            )
                        else:
                            # Shift (0 common) or Ambiguity (>1 common or multiple suggestions with no overlap)
                            merged_clar = list(reference_scope)
                            for a in dispatch_agent_ids:
                                if a not in merged_clar:
                                    merged_clar.append(a)

                            if not common:
                                # Case 1: Agent Domain Shift
                                msg = (
                                    f"Your previous questions were handled by "
                                    f"{', '.join(_agent_display_labels(reference_scope))}. "
                                    "This question seems more related to "
                                    f"{', '.join(_agent_display_labels(dispatch_agent_ids))}. "
                                    "Which system should I use?"
                                )
                            else:
                                # Case 2: Broad Ambiguity
                                msg = (
                                    "Your question could apply to several agents including your previous one. "
                                    "Which system or dataset should we use?"
                                )

                            state.intent_routing = {
                                **ir,
                                "needs_clarification": True,
                                "clarification_agent_ids": merged_clar,
                                "clarification_message": msg,
                            }
                            dispatch_agent_ids = []
                    elif (fan_out_all or non_multi_too_many) and intent_str != "multi_agent":
                        # Ambiguity with no prior scope (first question)
                        state.intent_routing = {
                            **ir,
                            "needs_clarification": True,
                            "clarification_message": ir.get("clarification_message") or (
                                "Your question could apply to several agents. "
                                "Which system or dataset should we use?"
                            ),
                        }
                        dispatch_agent_ids = []
            elif scope_ids:
                dispatch_agent_ids = [str(scope_ids[0])]

        ir = state.intent_routing or {}
        if bool(ir.get("needs_clarification")):
            dispatch_agent_ids = []
        else:
            dispatch_agent_ids = _restrict_dispatch_agent_ids_for_file_datasource(
                state, dispatch_agent_ids, scope_ids
            )
        intent_value = str(ir.get("intent", "")).strip().lower()
        # Respect explicit client datasource selection. When unset, infer from dispatched agents'
        # enabled datasources first; if several remain, use IR intent (sql_query→DATABASE, api→API, …).
        explicit_ds = str(state.datasource or "").strip()
        if not explicit_ds:
            inferred = _infer_talkback_datasource_from_dispatch_and_intent(
                dispatch_agent_ids, intent_value
            )
            if inferred:
                setattr(state.request, "datasource", inferred)
                state.datasource = inferred

        # Narrow dispatch_agent_ids to only those that actually support the selected (or inferred) datasource.
        # When "All" agents are selected (large dispatch scope), this ensures we don't fan out to agents
        # that cannot possibly handle the request (e.g. intent=document but agent has only SQL).
        if state.datasource:
            dispatch_agent_ids = _filter_agent_ids_by_datasource(dispatch_agent_ids, state.datasource)

        state.dispatch_agent_ids = dispatch_agent_ids
        return state


LLM_PRIOR_DF_CONSENT_MESSAGE = (
    "Do you allow us to send your previous answer to the AI so it can work with that result for "
    "this request? If your request needs web search, the AI may also use external search as part "
    "of the answer. Please confirm to continue."
)
LLM_PRIOR_DF_MISSING_DATA_MESSAGE = (
    "We could not use your previous answer for this step. Please confirm that you allow the AI to "
    "use your previous answer, then try again."
)


class _PriorDfConsentNode(BaseNode):
    """Gate prior-table @askme / @web LLM follow-ups: consent + df when the question uses those tags.
    Explicit @askme/@web clears agent disambiguation so llm_prior_df consent runs before ClarificationShortCircuit.
    If routing alone picks askme/web without tags, prior_df_followup may be cleared for normal SQL/API routing.
    Post-consent continuation may omit tags for one turn when ``pending_prior_df_followup`` is set."""

    def __init__(self, get_service: GetAskmeService, run_in_threadpool: RunInThreadpool):
        self._get_service = get_service
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "PriorDfConsentNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        from app.core.config import settings

        if not state.service or not state.session_key:
            return state

        ir = state.intent_routing
        if not ir or not isinstance(ir, dict) or ir.get("error"):
            return state

        # @askme / @web: prior-table LLM path takes precedence over agent disambiguation — otherwise
        # needs_clarification short-circuits this node and ClarificationShortCircuit returns "which agent?"
        # instead of llm_prior_df consent.
        q_early = (state.question_text or "").strip().lower()
        has_web_tag = "@web" in q_early
        has_askme_tag = "@askme" in q_early and not has_web_tag
        if has_web_tag or has_askme_tag:
            _p = "web_enrichment" if has_web_tag else "askme_analysis"
            ir = {
                **ir,
                "needs_clarification": False,
                "clarification_agent_ids": [],
                "clarification_message": None,
                "prior_df_followup": _p,
                "misrouted_web": False,
            }
            state.intent_routing = ir
        elif ir.get("needs_clarification"):
            return state

        # Clear stale llm_prior_df pending if routing no longer requests prior-df follow-up
        if state.session_state and state.session_key:
            try:
                cc_raw = state.session_state.get("clarification_context_json")
                cc_prev = json.loads(cc_raw) if cc_raw else {}
            except Exception:
                cc_prev = {}
            if isinstance(cc_prev, dict) and cc_prev.get("consent_type") == "llm_prior_df":
                prior0 = str(ir.get("prior_df_followup", "none")).lower()
                pp0 = getattr(state, "pending_prior_df_followup", None)
                effective0 = pp0 if pp0 in ("askme_analysis", "web_enrichment") else prior0
                if effective0 not in ("askme_analysis", "web_enrichment"):
                    await self._run_in_threadpool(
                        state.service.clear_talkback_pending_clarification_only, state.session_key
                    )
                    state.session_state = await self._run_in_threadpool(
                        state.service.get_talkback_session_state, state.session_key
                    )

        prior = str(ir.get("prior_df_followup", "none")).lower()
        pp = getattr(state, "pending_prior_df_followup", None)
        if pp in ("askme_analysis", "web_enrichment"):
            prior = pp
            ir = {**ir, "prior_df_followup": prior}
            state.intent_routing = ir

        q_lower = (state.question_text or "").strip().lower()
        has_web = "@web" in q_lower
        has_askme_only = "@askme" in q_lower and not has_web

        if has_web:
            prior = "web_enrichment"
            ir = {**ir, "prior_df_followup": "web_enrichment"}
            state.intent_routing = ir
        elif has_askme_only:
            prior = "askme_analysis"
            ir = {**ir, "prior_df_followup": "askme_analysis"}
            state.intent_routing = ir
        elif prior in ("askme_analysis", "web_enrichment"):
            # Router-only mode without tags: normal SQL/API/multi-agent talkback, not llm_prior_df.
            # Exception: user just confirmed consent (`pp`); their next message may omit tags once.
            if pp in ("askme_analysis", "web_enrichment"):
                pass
            else:
                state.intent_routing = {**ir, "prior_df_followup": "none"}
                return state

        if prior not in ("askme_analysis", "web_enrichment"):
            return state
        if ir.get("misrouted_web"):
            return state
        # Consent does not require server-side history: after context clear the client may still
        # hold the prior table and will send it in `df` after opt-in.

        if prior == "askme_analysis" and not settings.IS_ASKME_ACTIVE:
            state.response = {
                "success": False,
                "message": "@askme analysis is disabled for this deployment.",
                "details": {"feature": "IS_ASKME_ACTIVE"},
                "intent_routing": state.intent_routing,
                "context_cleared": state.context_cleared,
                "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()
        if prior == "web_enrichment" and not settings.IS_WEB_ACTIVE:
            state.response = {
                "success": False,
                "message": "@web search is disabled for this deployment.",
                "details": {"feature": "IS_WEB_ACTIVE"},
                "intent_routing": state.intent_routing,
                "context_cleared": state.context_cleared,
                "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()

        search_web = bool(getattr(state.request, "search_web", False))

        if not search_web:
            service_key = state.resolved_agent_ids if len(state.resolved_agent_ids) > 1 else state.resolved_agent_ids[0]
            service = self._get_service(agent_id=service_key)
            if service:
                sk = _derive_talkback_db_session_key(service, state)
                dispatch_scope = list(state.dispatch_agent_ids) if state.dispatch_agent_ids else list(state.resolved_agent_ids)
                consent_rid = _derive_request_id(state)
                consent_meta = {
                    "turn_type": "clarification",
                    "consent_type": "llm_prior_df",
                    "prior_df_followup": prior,
                }
                await self._run_in_threadpool(
                    service.append_turn_to_talkback_session,
                    sk,
                    state.current_user_id,
                    state.resolved_agent_ids,
                    state.question_text,
                    LLM_PRIOR_DF_CONSENT_MESSAGE,
                    10,
                    consent_rid,
                    consent_meta,
                    dispatch_scope,
                )
                await self._run_in_threadpool(
                    service.upsert_talkback_session_state,
                    sk,
                    state.current_user_id,
                    state.resolved_agent_ids,
                    None,
                    None,
                    True,
                    LLM_PRIOR_DF_CONSENT_MESSAGE,
                    {"consent_type": "llm_prior_df", "prior_df_followup": prior},
                    dispatch_scope,
                )

            state.response = {
                "success": True,
                "message": "Consent required",
                "details": {
                    "needs_clarification": True,
                    "clarification_message": LLM_PRIOR_DF_CONSENT_MESSAGE,
                    "consent_type": "llm_prior_df",
                },
                "intent_routing": state.intent_routing,
                "context_cleared": state.context_cleared,
                "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()

        raw_df = getattr(state.request, "df", None)
        rows: Optional[List[Dict[str, Any]]] = None
        if raw_df is not None:
            rows = await self._run_in_threadpool(state.service._parse_df_input, raw_df)
            if isinstance(rows, list) and len(rows) == 0:
                rows = None

        if not rows:
            state.response = {
                "success": False,
                "message": LLM_PRIOR_DF_MISSING_DATA_MESSAGE,
                "details": {"reason": "missing_prior_df"},
                "intent_routing": state.intent_routing,
                "context_cleared": state.context_cleared,
                "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()

        full_raw = (state.question_text or "").strip()
        eff_prior = prior
        if eff_prior == "web_enrichment" and not settings.IS_WEB_ACTIVE:
            state.response = {
                "success": False,
                "message": "@web search is disabled for this deployment.",
                "details": {"feature": "IS_WEB_ACTIVE"},
                "intent_routing": state.intent_routing,
                "context_cleared": state.context_cleared,
                "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()
        if eff_prior == "askme_analysis" and not settings.IS_ASKME_ACTIVE:
            state.response = {
                "success": False,
                "message": "@askme analysis is disabled for this deployment.",
                "details": {"feature": "IS_ASKME_ACTIVE"},
                "intent_routing": state.intent_routing,
                "context_cleared": state.context_cleared,
                "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()

        state.question_text = full_raw
        setattr(state.request, "question", full_raw)
        setattr(state.request, "df", rows)

        return state


class _BroadScopeProcessingWarningNode(BaseNode):
    """
    Before downstream agent clarification / execute: if intent classification indicates the query is
    broadly scoped (LLM), offer a single-step talkback prompt. Follow-up is ``proceed`` or a
    time-only fragment merged with the stored original question in ``LoadSessionNode``.
    """

    def __init__(self, get_service: GetAskmeService, run_in_threadpool: RunInThreadpool):
        self._get_service = get_service
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "BroadScopeProcessingWarningNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        if getattr(state, "skip_broad_scope_check_once", False):
            state.skip_broad_scope_check_once = False
            return state
        if getattr(state, "agent_disambiguation_follow_up", False):
            return state
        ir = state.intent_routing
        if not ir or not isinstance(ir, dict) or ir.get("error"):
            return state
        intent = str(ir.get("intent", "")).strip().lower()
        if intent == "action":
            return state
        prior = str(ir.get("prior_df_followup", "none")).lower()
        if prior in ("askme_analysis", "web_enrichment", PRIOR_DF_FOLLOWUP_FILE_RECONCILE):
            return state

        q = (state.question_text or "").strip()
        if not q:
            return state
        # Short noun-like text matches agent-pick replies; only skip broad-scope when that is actually expected.
        if _looks_like_clarification_reply(q) and _history_last_turn_expected_agent_disambiguation_reply(state):
            return state

        service_key = (
            state.resolved_agent_ids if len(state.resolved_agent_ids) > 1 else state.resolved_agent_ids[0]
        )
        service = self._get_service(agent_id=service_key)
        if not service:
            return state

        def _broad_check() -> bool:
            return _talkback_is_broad_or_underconstrained_query(
                service,
                q,
                temperature=0.0,
                max_tokens=80,
                use_session=False,
            )

        is_broad = await self._run_in_threadpool(_broad_check)
        if not is_broad:
            return state

        request_id = _derive_request_id(state)
        session_key = _derive_talkback_db_session_key(service, state)
        dispatch_scope = (
            list(state.dispatch_agent_ids) if state.dispatch_agent_ids else list(state.resolved_agent_ids)
        )
        clar_meta: Dict[str, Any] = {
            "turn_type": "clarification",
            "needs_clarification": True,
            "consent_type": BROAD_SCOPE_PROCESSING_CONSENT_TYPE,
        }
        if ir.get("intent"):
            clar_meta["intent"] = str(ir.get("intent"))[:120]
        await self._run_in_threadpool(
            service.append_turn_to_talkback_session,
            session_key,
            state.current_user_id,
            state.resolved_agent_ids,
            q,
            BROAD_SCOPE_CLARIFICATION_MESSAGE,
            10,
            request_id,
            clar_meta,
            dispatch_scope,
        )
        cc_ctx: Dict[str, Any] = {
            **ir,
            "consent_type": BROAD_SCOPE_PROCESSING_CONSENT_TYPE,
            "captured_question": q,
        }
        if isinstance(state.file_cmd_parse, dict):
            cc_ctx["file_cmd_parse"] = state.file_cmd_parse
        await self._run_in_threadpool(
            service.upsert_talkback_session_state,
            session_key,
            state.current_user_id,
            state.resolved_agent_ids,
            None,
            None,
            True,
            BROAD_SCOPE_CLARIFICATION_MESSAGE,
            cc_ctx,
            dispatch_scope,
        )

        state.response = {
            "success": True,
            "message": "Scope clarification",
            "details": {
                "needs_clarification": True,
                "clarification_message": BROAD_SCOPE_CLARIFICATION_MESSAGE,
                "consent_type": BROAD_SCOPE_PROCESSING_CONSENT_TYPE,
                "clarification_agent_ids": [],
            },
            "intent_routing": ir,
            "context_cleared": state.context_cleared,
            "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
            "trace_id": state.trace.trace_id,
        }
        raise _TalkbackStop()


class _ClarificationShortCircuitNode(BaseNode):
    def __init__(self, get_service: GetAskmeService, run_in_threadpool: RunInThreadpool):
        self._get_service = get_service
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "ClarificationShortCircuitNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        ir = state.intent_routing
        # Prior table vs spreadsheet upload: steer to reconcile (HTTP or talkback multipart), not FILE/@file routing.
        if ir and isinstance(ir, dict) and not ir.get("error"):
            prior = str(ir.get("prior_df_followup", "none")).lower()
            if prior == PRIOR_DF_FOLLOWUP_FILE_RECONCILE:
                clar_msg = (
                    str(ir.get("clarification_message") or "").strip()
                    or DEFAULT_FILE_RECONCILE_CLARIFICATION_MESSAGE
                )
                service_key = (
                    state.resolved_agent_ids
                    if len(state.resolved_agent_ids) > 1
                    else state.resolved_agent_ids[0]
                )
                service = self._get_service(agent_id=service_key)
                if service:
                    request_id = _derive_request_id(state)
                    session_key = _derive_talkback_db_session_key(service, state)
                    dispatch_scope = (
                        list(state.dispatch_agent_ids)
                        if state.dispatch_agent_ids
                        else list(state.resolved_agent_ids)
                    )
                    clar_meta: Dict[str, Any] = {
                        "turn_type": "clarification",
                        "needs_clarification": True,
                        "consent_type": "file_reconcile",
                        "prior_df_followup": PRIOR_DF_FOLLOWUP_FILE_RECONCILE,
                    }
                    if ir.get("intent"):
                        clar_meta["intent"] = str(ir.get("intent"))[:120]
                    await self._run_in_threadpool(
                        service.append_turn_to_talkback_session,
                        session_key,
                        state.current_user_id,
                        state.resolved_agent_ids,
                        state.question_text,
                        clar_msg,
                        10,
                        request_id,
                        clar_meta,
                        dispatch_scope,
                    )
                    await self._run_in_threadpool(
                        service.upsert_talkback_session_state,
                        session_key,
                        state.current_user_id,
                        state.resolved_agent_ids,
                        None,
                        None,
                        True,
                        clar_msg,
                        {
                            **ir,
                            "prior_df_followup": PRIOR_DF_FOLLOWUP_FILE_RECONCILE,
                            "file_cmd_parse": state.file_cmd_parse if isinstance(state.file_cmd_parse, dict) else None,
                        },
                        dispatch_scope,
                    )
                state.response = {
                    "success": True,
                    "message": "Reconcile upload required",
                    "details": {
                        "needs_clarification": True,
                        "clarification_message": clar_msg,
                        "consent_type": "file_reconcile",
                        "clarification_agent_ids": [],
                        "suggested_questions": [],
                    },
                    "intent_routing": ir,
                    "context_cleared": state.context_cleared,
                    "context_clear_reason": state.context_clear_reason
                    if state.context_cleared
                    else None,
                    "trace_id": state.trace.trace_id,
                }
                raise _TalkbackStop()

        if not (ir and isinstance(ir, dict) and not ir.get("error") and ir.get("needs_clarification")):
            return state

        service_key = state.resolved_agent_ids if len(state.resolved_agent_ids) > 1 else state.resolved_agent_ids[0]
        service = self._get_service(agent_id=service_key)

        cids = ir.get("clarification_agent_ids")
        if not isinstance(cids, list):
            cids = []
        clar_agent_ids = [str(x) for x in cids if str(x).strip()]

        # Suggest starter questions only for generic/out-of-scope clarifications.
        # For agent disambiguation prompts (clarification_agent_ids present), do not
        # include suggested_questions in the response payload.
        suggestion_agent_ids = (
            []
            if clar_agent_ids
            else [str(x) for x in state.resolved_agent_ids if str(x).strip()]
        )
        suggested_questions: List[str] = []
        if suggestion_agent_ids:
            per_agent_limit = 5 if clar_agent_ids else 1
            for aid in suggestion_agent_ids:
                try:
                    aid_service = self._get_service(agent_id=aid)
                    if not aid_service:
                        continue
                    qs, _ = await self._run_in_threadpool(
                        aid_service.get_agent_questions_from_db,
                        aid,
                        per_agent_limit,
                    )
                    qlist = [str(q).strip() for q in (qs or []) if str(q).strip()]
                    if not qlist:
                        continue
                    suggested_questions.extend(qlist)
                except Exception:
                    continue

        # De-duplicate while preserving order.
        if suggested_questions:
            seen = set()
            deduped: List[str] = []
            for q in suggested_questions:
                k = q.lower()
                if k in seen:
                    continue
                seen.add(k)
                deduped.append(q)
            suggested_questions = deduped[:5]

        base_clar_msg = str(ir.get("clarification_message") or "").strip()
        is_agent_disambiguation = bool(clar_agent_ids)
        is_custom_file_clar = bool(ir.get("is_custom_file_clarification"))
        if suggested_questions and not is_agent_disambiguation and not is_custom_file_clar:
            guidance = "This looks out of scope for the current agents. Would you like to ask one of these instead?"
            if base_clar_msg:
                clar_msg = f"{base_clar_msg}\n\n{guidance}"
            else:
                clar_msg = guidance
        else:
            clar_msg = base_clar_msg

        if service:
            request_id = _derive_request_id(state)
            session_key = _derive_talkback_db_session_key(service, state)
            dispatch_scope = list(state.dispatch_agent_ids) if state.dispatch_agent_ids else list(state.resolved_agent_ids)
            clar_meta: Dict[str, Any] = {
                "turn_type": "clarification",
                "needs_clarification": True,
            }
            if ir.get("intent"):
                clar_meta["intent"] = str(ir.get("intent"))[:120]
            if clar_agent_ids:
                clar_meta["clarification_agent_ids"] = clar_agent_ids
            await self._run_in_threadpool(
                service.append_turn_to_talkback_session,
                session_key,
                state.current_user_id,
                state.resolved_agent_ids,
                state.question_text,
                clar_msg,
                10,
                request_id,
                clar_meta,
                dispatch_scope,
            )
            await self._run_in_threadpool(
                service.upsert_talkback_session_state,
                session_key,
                state.current_user_id,
                state.resolved_agent_ids,
                None,
                None,
                True,
                clar_msg or "Could you clarify your question?",
                {
                    **ir,
                    "file_cmd_parse": state.file_cmd_parse if isinstance(state.file_cmd_parse, dict) else None,
                },
                dispatch_scope,
            )

        details: Dict[str, Any] = {
            "needs_clarification": True,
            "clarification_message": clar_msg,
        }
        if clar_agent_ids and not is_custom_file_clar:
            details["clarification_agent_ids"] = clar_agent_ids
        for k in ("consent_type", "clarification_type", "clarification_options"):
            if k in ir:
                details[k] = ir[k]

        if suggested_questions and not is_agent_disambiguation and not is_custom_file_clar:
            details["suggested_questions"] = suggested_questions

        state.response = {
            "success": True,
            "message": "Clarification needed",
            "details": details,
            "intent_routing": ir,
            "context_cleared": state.context_cleared,
            "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
            "trace_id": state.trace.trace_id,
        }
        raise _TalkbackStop()


class _ExecuteNode(BaseNode):
    def __init__(self, execute: ExecuteTalkback):
        self._execute = execute

    @property
    def name(self) -> str:
        return "ExecuteNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        ir_exec = state.intent_routing
        if ir_exec and isinstance(ir_exec, dict) and not ir_exec.get("error") and ir_exec.get(
            "needs_clarification"
        ):
            details = {
                "needs_clarification": True,
                "clarification_message": (
                    str(ir_exec.get("clarification_message") or "").strip()
                    or "Could you clarify which system or dataset should be used?"
                ),
            }
            c_ids = ir_exec.get("clarification_agent_ids")
            if c_ids and not ir_exec.get("is_custom_file_clarification"):
                details["clarification_agent_ids"] = [
                    str(x)
                    for x in (c_ids or [])
                    if str(x).strip()
                ]
            for k in ("consent_type", "clarification_type", "clarification_options"):
                if k in ir_exec:
                    details[k] = ir_exec[k]

            state.response = {
                "success": True,
                "message": "Clarification needed",
                "details": details,
                "intent_routing": ir_exec,
                "context_cleared": state.context_cleared,
                "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()

        if not state.dispatch_agent_ids:
            state.response = {
                "success": False,
                "message": "No agents available for dispatch",
                "details": {},
                "intent_routing": state.intent_routing,
                "context_cleared": state.context_cleared,
                "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                "trace_id": state.trace.trace_id,
            }
            raise _TalkbackStop()

        ir_exec = state.intent_routing
        if ir_exec and isinstance(ir_exec, dict) and not ir_exec.get("error"):
            if (
                str(ir_exec.get("prior_df_followup", "none")).lower()
                == PRIOR_DF_FOLLOWUP_FILE_RECONCILE
            ):
                state.response = {
                    "success": True,
                    "message": "Reconcile upload required",
                    "details": {
                        "needs_clarification": True,
                        "clarification_message": DEFAULT_FILE_RECONCILE_CLARIFICATION_MESSAGE,
                        "consent_type": "file_reconcile",
                    },
                    "intent_routing": ir_exec,
                    "context_cleared": state.context_cleared,
                    "context_clear_reason": state.context_clear_reason
                    if state.context_cleared
                    else None,
                    "trace_id": state.trace.trace_id,
                }
                raise _TalkbackStop()

        if state.file_cmd_parse:
            from app.routes.askme import _ask_file_core

            merged_ids: List[str] = []
            seen_ids = set()
            for lst in (state.resolved_agent_ids or [], state.dispatch_agent_ids or []):
                for aid in lst:
                    s = str(aid).strip()
                    if s and s not in seen_ids:
                        seen_ids.add(s)
                        merged_ids.append(s)
            file_agent_id = _resolve_talkback_file_agent_id(state, merged_ids)
            if not file_agent_id:
                state.response = {
                    "success": False,
                    "message": "No file-capable agents are available for this request.",
                    "details": {
                        "needs_clarification": True,
                        "clarification_message": (
                            "Please select an agent that supports file/document datasource."
                        ),
                        "error_code": "file_agent_required",
                    },
                    "intent_routing": state.intent_routing,
                    "context_cleared": state.context_cleared,
                    "context_clear_reason": state.context_clear_reason
                    if state.context_cleared
                    else None,
                    "trace_id": state.trace.trace_id,
                }
                raise _TalkbackStop()

            async def _run_non_file() -> Any:
                req_non_file = state.request
                non_file_clause = str((state.file_cmd_parse or {}).get("non_file_clause") or "").strip()
                if not non_file_clause:
                    # If the entire query was just a file command (e.g. "@file find policy"),
                    # there is no separate database query to execute.
                    return None

                req_non_file = _clone_ask_request(
                    state.request,
                    {
                        "question": non_file_clause,
                        "datasource": None,
                    },
                )
                _ensure_web_prefix_for_search_web_request(req_non_file, state)
                setattr(
                    req_non_file,
                    "agent_id",
                    list(
                        getattr(state, "_non_file_dispatch_agent_ids_for_mixed", None)
                        or state.dispatch_agent_ids
                    ),
                )
                return await self._execute(req_non_file, state.current_user)

            async def _run_file() -> Any:
                return await _ask_file_core(
                    user_id=state.current_user_id,
                    agent_id=file_agent_id,
                    question=(state.file_cmd_parse or {}).get("file_clause") or "",
                    current_user=state.current_user,
                    skip_log=False,
                )

            # If there is no non-file clause, only execute the file agent and return its result directly.
            non_file_clause = str((state.file_cmd_parse or {}).get("non_file_clause") or "").strip()
            if not non_file_clause:
                file_raw = await _run_file()
                state.result = _coerce_route_handler_result_to_dict(file_raw)
                state.result["intent_routing"] = state.intent_routing
                state.result["context_cleared"] = state.context_cleared
                if state.context_cleared:
                    state.result["context_clear_reason"] = state.context_clear_reason
                state.result["trace_id"] = state.trace.trace_id
                _update_dispatch_from_result(state)
                return state

            non_file_raw, file_raw = await asyncio.gather(
                _run_non_file(), _run_file(), return_exceptions=True
            )
            state.result = _merge_mixed_intent_results(non_file_raw, file_raw, state)
            return state

        logger.info(
            "TalkbackOrchestrator trace_id=%s node=%s executing agent_ids=%s",
            state.trace.trace_id,
            self.name,
            state.dispatch_agent_ids,
        )
        _ensure_web_prefix_for_search_web_request(state.request, state)
        setattr(state.request, "agent_id", list(state.dispatch_agent_ids))
        raw = await self._execute(state.request, state.current_user)

        state.result = _coerce_route_handler_result_to_dict(raw)
        state.result["intent_routing"] = state.intent_routing
        state.result["context_cleared"] = state.context_cleared
        if state.context_cleared:
            state.result["context_clear_reason"] = state.context_clear_reason
        state.result["trace_id"] = state.trace.trace_id
        _update_dispatch_from_result(state)
        return state


def _merge_mixed_intent_results(
    non_file_raw: Any, file_raw: Any, state: TalkbackState
) -> Dict[str, Any]:
    from app.services.talkback.utils import _infer_datasource_from_details

    if isinstance(non_file_raw, Exception):
        non_file_raw = {"success": False, "message": str(non_file_raw), "details": {}}
    if isinstance(file_raw, Exception):
        file_raw = {"success": False, "message": str(file_raw), "details": {}}

    non_file_payload = _coerce_route_handler_result_to_dict(non_file_raw)
    file_payload = _coerce_route_handler_result_to_dict(file_raw)
    non_file_ok = bool(non_file_payload.get("success"))
    file_ok = bool(file_payload.get("success"))

    non_file_details = (
        non_file_payload.get("details")
        if isinstance(non_file_payload.get("details"), dict)
        else {}
    )
    file_details = (
        file_payload.get("details")
        if isinstance(file_payload.get("details"), dict)
        else {}
    )
    non_file_datasource = _infer_datasource_from_details(non_file_details) or "unknown"

    warnings: List[str] = []
    if not non_file_ok:
        warnings.append(
            f"{non_file_datasource} branch failed: {str(non_file_payload.get('message') or '')}"
        )

    # If file search matched specific agents, refine the dispatch scope
    matched = file_details.get("matched_agent_ids")
    if matched and isinstance(matched, list):
        new_ids = [str(a) for a in matched if a]
        if _question_requires_file_only_dispatch(state):
            new_ids = _filter_agent_ids_by_datasource(new_ids, "File")
        if new_ids:
            state.dispatch_agent_ids = new_ids
            logger.info(
                "TalkbackState refined dispatch_agent_ids to %s based on file branch results",
                state.dispatch_agent_ids,
            )
    if not file_ok:
        warnings.append(f"File branch failed: {str(file_payload.get('message') or '')}")
    partial = not (non_file_ok and file_ok)

    merged_details = {
        **non_file_details,
        "file_result": file_details,
        "partial": partial,
        "branch_status": {
            "non_file": "success" if non_file_ok else "failed",
            "non_file_datasource": non_file_datasource,
            "file": "success" if file_ok else "failed",
        },
        "warnings": warnings,
    }

    log_meta = json.dumps(
        {
            "mixed_intent": True,
            "non_file_status": merged_details["branch_status"]["non_file"],
            "non_file_datasource": non_file_datasource,
            "file_status": merged_details["branch_status"]["file"],
            "partial": partial,
            "warnings": warnings,
        },
        default=str,
    )

    out_merged: Dict[str, Any] = {
        "success": non_file_ok or file_ok,
        "message": (
            "Question processed successfully (mixed intent)"
            if (non_file_ok or file_ok)
            else "Both branches failed"
        ),
        "details": merged_details,
        "intent_routing": state.intent_routing,
        "context_cleared": state.context_cleared,
        "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
        "trace_id": state.trace.trace_id,
        "_log_meta": log_meta,
    }
    if state.dispatch_agent_ids:
        out_merged["dispatch_agent_ids"] = list(state.dispatch_agent_ids)
    elif _question_requires_file_only_dispatch(state):
        fo_ids = _filter_agent_ids_by_datasource(list(state.resolved_agent_ids or []), "File")
        if fo_ids:
            out_merged["dispatch_agent_ids"] = list(fo_ids)
    return out_merged


class _PersistTurnNode(BaseNode):
    def __init__(self, get_service: GetAskmeService, run_in_threadpool: RunInThreadpool):
        self._get_service = get_service
        self._run_in_threadpool = run_in_threadpool

    @property
    def name(self) -> str:
        return "PersistTurnNode"

    async def process(self, state: TalkbackState) -> TalkbackState:
        if not isinstance(state.result, dict):
            state.response = {
                "success": False,
                "message": "Talkback execution returned an unexpected response shape.",
                "details": {"reason": "non_dict_result"},
                "intent_routing": state.intent_routing,
                "context_cleared": state.context_cleared,
                "context_clear_reason": state.context_clear_reason if state.context_cleared else None,
                "trace_id": state.trace.trace_id,
            }
            return state

        persist_dispatch_ids: List[str] = [str(x) for x in (state.dispatch_agent_ids or []) if str(x).strip()]
        if not persist_dispatch_ids and _question_requires_file_only_dispatch(state):
            persist_dispatch_ids = _filter_agent_ids_by_datasource(
                list(state.resolved_agent_ids or []), "File"
            )
        if not persist_dispatch_ids:
            if _question_requires_file_only_dispatch(state):
                if isinstance(state.result, dict):
                    state.result["dispatch_agent_ids"] = []
                state.response = state.result
                return state
            persist_dispatch_ids = [str(x) for x in (state.resolved_agent_ids or []) if str(x).strip()]
        if not persist_dispatch_ids:
            state.response = state.result
            return state

        state.dispatch_agent_ids = list(persist_dispatch_ids)

        service_key = persist_dispatch_ids if len(persist_dispatch_ids) > 1 else persist_dispatch_ids[0]
        service = self._get_service(agent_id=service_key)
        if not service:
            # Still return the handler result; only session persistence is skipped.
            state.response = state.result
            return state

        # Same scope as LoadSessionNode / clarification: resolved_agent_ids, not dispatch-only.
        # Otherwise history is stored under one key and loaded under another when routing narrows to one agent.
        session_key = state.session_key or _derive_talkback_db_session_key(service, state)
        answer_text = str(state.result.get("message", "") or "")
        clar_sel = getattr(state, "disambiguation_user_choice", None)
        answer_meta = build_talkback_answer_meta(
            state.result,
            dispatch_agent_ids=list(persist_dispatch_ids),
            clarification_selection=clar_sel if isinstance(clar_sel, str) else None,
        )
        request_id = _derive_request_id(state)
        raw_hist = (getattr(state, "talkback_raw_question_text", None) or "").strip()
        details_q = ""
        try:
            details_obj = state.result.get("details") if isinstance(state.result, dict) else None
            if isinstance(details_obj, dict):
                details_q = str(details_obj.get("question", "") or "").strip()
        except Exception:
            details_q = ""
        question_for_history = (
            getattr(state, "talkback_history_question_text", None)
            or (details_q if details_q else None)
            or (raw_hist if raw_hist else None)
            or state.question_text
        )
        _print_node(
            "call",
            self.name,
            state.trace.trace_id,
            state,
            fn="append_turn_to_talkback_session",
            session_key=_safe_str(session_key, 96),
            agent_ids=list(persist_dispatch_ids),
            question=_safe_str(question_for_history),
            answer_preview=_safe_str(answer_text),
            limit=10,
            request_id=_safe_str(request_id, 96),
        )
        dispatch_scope = list(persist_dispatch_ids)
        await self._run_in_threadpool(
            service.append_turn_to_talkback_session,
            session_key,
            state.current_user_id,
            dispatch_scope,
            question_for_history,
            answer_text,
            10,
            request_id,
            answer_meta,
            dispatch_scope,
        )

        try:
            agent = service._ensure_agent()
            last_id = getattr(agent, "_last_session_response_id", None)
        except Exception:
            last_id = None

        if last_id:
            await self._run_in_threadpool(
                service.upsert_talkback_session_state,
                session_key,
                state.current_user_id,
                dispatch_scope,
                None,
                last_id,
                False,
                None,
                {},
                dispatch_scope,
            )
        # Only mark the final response after persistence.
        if isinstance(state.result, dict):
            state.result["dispatch_agent_ids"] = list(state.dispatch_agent_ids)
        state.response = state.result
        return state

def _update_dispatch_from_result(state: TalkbackState):
    """Refine dispatch_agent_ids if the result contains specifically matched agents."""
    if not state.result or not isinstance(state.result.get("details"), dict):
        return
    
    matched = state.result["details"].get("matched_agent_ids")
    if matched and isinstance(matched, list):
        # We found specific agents that have the files. 
        # Update the dispatch scope so subsequent talkback/follow-ups 
        # are pinned to these specific agents.
        new_ids = [str(a) for a in matched if a]
        if _question_requires_file_only_dispatch(state):
            new_ids = _filter_agent_ids_by_datasource(new_ids, "File")
        if new_ids:
            state.dispatch_agent_ids = new_ids
            logger.info(
                "TalkbackState refined dispatch_agent_ids to %s based on search results (from %s)",
                state.dispatch_agent_ids,
                state.result.get("result_type", "unknown")
            )
        else:
            logger.info("TalkbackState matched_agent_ids empty, keeping dispatch_agent_ids %s", state.dispatch_agent_ids)
    elif _question_requires_file_only_dispatch(state):
        fo = _filter_agent_ids_by_datasource(list(state.resolved_agent_ids or []), "File")
        if fo:
            state.dispatch_agent_ids = list(fo)
            logger.info(
                "TalkbackState set dispatch_agent_ids to file-capable agents %s (no matched_agent_ids in result)",
                state.dispatch_agent_ids,
            )
    else:
        logger.info("TalkbackState no matched_agent_ids in result details, keeping dispatch_agent_ids %s", state.dispatch_agent_ids)

def _is_file_search_clarification_response(state: TalkbackState) -> bool:
    """Check if the current question is a response to the 'exact vs similar' file search clarification."""
    if not state.history_turns:
        return False
    
    last = state.history_turns[-1]
    if not isinstance(last, dict):
        return False
    
    _meta = last.get("answer_meta")
    if not isinstance(_meta, dict) or not _meta.get("needs_clarification"):
        return False
    
    _last_ans = str(last.get("answer") or "").strip().lower()
    if not ("find the exact file" in _last_ans or "find similar files" in _last_ans or "select one option to continue" in _last_ans):
        return False
    
    _q_lower = (state.question_text or "").strip().lower()
    # Check if the question matches one of the labels or actions
    return "exact" in _q_lower or "similar" in _q_lower or "first" in _q_lower or "second" in _q_lower
