"""
Training-driven multi-agent pipeline: load agent context → decompose question →
map sub-questions to agents/datasources → execute → merge tabular results.

File/document sub-tasks return structured payloads (not merged with SQL/API tables).
Falls back to legacy routing when this module returns used_pipeline=False.
"""
from __future__ import annotations

import asyncio
import json
import re
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, List, Optional, Tuple


import pandas as pd

# ---------------------------------------------------------------------------
# Tiktoken – optional token counter (OpenAI tokeniser)
# ---------------------------------------------------------------------------
try:
    import tiktoken as _tiktoken  # type: ignore[import]
    _TIKTOKEN_AVAILABLE = True
except ImportError:
    _tiktoken = None  # type: ignore[assignment]
    _TIKTOKEN_AVAILABLE = False

_TIKTOKEN_DEFAULT_MODEL = "gpt-4.1"
_TIKTOKEN_FALLBACK_ENCODING = "cl100k_base"
# Encoder cache – pay the load cost only once per model name.
_tiktoken_encoder_cache: Dict[str, Any] = {}

# Max chars of training/DDL text per agent sent to the planner LLM (token safety)
_TRAINING_SNIP_LEN = 6000

# Default threshold for planning-time filtering.
# Sub-tasks with confidence_score < this value are not executed.
_DEFAULT_THRESHOLD = 0.6

# Default threshold for result-time filtering (before merging).
# 0.0 keeps all non-empty tabular results.
_RESULT_SCORE_THRESHOLD = 0.0


def _all_frames_disjoint(frames: List[pd.DataFrame]) -> bool:
    """Return True when every frame pair has zero business-column overlap."""
    meta_cols = {"_datasource_type", "_agents"}
    col_sets = [{str(c) for c in df.columns if str(c) not in meta_cols} for df in frames]
    for i in range(len(col_sets)):
        for j in range(i + 1, len(col_sets)):
            if col_sets[i] & col_sets[j]:
                return False
    return True


def _log(stage: str, message: str, **meta: Any) -> None:
    """Small structured logger for pipeline debugging."""
    details = " ".join(f"{k}={meta[k]!r}" for k in sorted(meta))
    if details:
        print(f"[multi_agent_pipeline] {stage} {message} | {details}")
    else:
        print(f"[multi_agent_pipeline] {stage} {message}")


# ---------------------------------------------------------------------------
# Token calculator (Tiktoken / OpenAI)
# ---------------------------------------------------------------------------

def _get_encoder(model: str = _TIKTOKEN_DEFAULT_MODEL) -> Any:
    """Return a cached tiktoken encoder for *model*, falling back to cl100k_base."""
    if not _TIKTOKEN_AVAILABLE:
        return None
    if model in _tiktoken_encoder_cache:
        return _tiktoken_encoder_cache[model]
    try:
        enc = _tiktoken.encoding_for_model(model)  # type: ignore[union-attr]
    except KeyError:
        enc = _tiktoken.get_encoding(_TIKTOKEN_FALLBACK_ENCODING)  # type: ignore[union-attr]
    _tiktoken_encoder_cache[model] = enc
    return enc


def _count_tokens(text: str, model: str = _TIKTOKEN_DEFAULT_MODEL) -> int:
    """Return the token count for *text* using tiktoken.

    Falls back to a rough word-based estimate when tiktoken is not installed.
    """
    if not text:
        return 0
    enc = _get_encoder(model)
    if enc is None:
        return max(1, len(text.split()))
    return len(enc.encode(text))


def _log_prompt_tokens(
    stage: str,
    sys_msg: str,
    user_msg: str,
    model: str = _TIKTOKEN_DEFAULT_MODEL,
) -> Dict[str, int]:
    """Count tokens for a system + user prompt pair, emit an emoji log line,
    and return the breakdown so callers can accumulate totals.

    Emoji guide
    -----------
    🪙  per-message token count header
    📏  system-message token count
    📝  user-message token count
    📊  combined total
    ⚠️  tiktoken unavailable – using word-count estimate
    """
    sys_tokens = _count_tokens(sys_msg, model)
    user_tokens = _count_tokens(user_msg, model)
    total = sys_tokens + user_tokens
    backend = "tiktoken" if _TIKTOKEN_AVAILABLE else "word-estimate"

    _log("🪙", f"token_calc:{stage}:begin", model=model, backend=backend)
    _log("📏", f"token_calc:{stage}:system_msg",  tokens=sys_tokens,  chars=len(sys_msg))
    _log("📝", f"token_calc:{stage}:user_msg",    tokens=user_tokens, chars=len(user_msg))
    _log("📊", f"token_calc:{stage}:total_input", tokens=total)
    if not _TIKTOKEN_AVAILABLE:
        _log("⚠️", "token_calc:tiktoken_unavailable_using_word_estimate")

    return {"sys_tokens": sys_tokens, "user_tokens": user_tokens, "total_tokens": total}


def load_training_data(agent_id: str) -> str:
    """Return init training data (DDL / docs) for an agent."""
    from app.services.agent_config import get_agent_config

    cfg = get_agent_config(agent_id=agent_id)
    if not cfg:
        return ""
    return (cfg.get("traning_data") or cfg.get("training_data") or "").strip()


def load_api_training_data(agent_id: str) -> List[Dict[str, Any]]:
    """Return API training rows (endpoint, base_params, description) for an agent."""
    try:
        from app.services.api_call import APIChatService
        from app.core.config import settings

        svc = APIChatService(
            settings.OPENAI_API_KEY,
            claude_api_key=(getattr(settings, "CLAUDE_API_KEY", None) or ""),
        )
        rows = svc.get_all_apis(agent_id)
        _log("📡", "load_api_training_data:loaded", agent_id=agent_id, count=len(rows or []))
        return rows or []
    except Exception as exc:
        _log("⚠️", "load_api_training_data:error", agent_id=agent_id, error=str(exc))
        return []


def build_agents_metadata(agent_ids: List[str]) -> List[Dict[str, Any]]:
    """Per-agent record: id, name, training excerpt, allowed datasources."""
    from app.services.agent_config import get_agent_by_id, ordered_available_datasources

    _log("🔎", "build_agents_metadata:start", agent_ids=agent_ids, count=len(agent_ids))
    out: List[Dict[str, Any]] = []
    for aid in agent_ids:
        _log("🧩", "build_agents_metadata:agent_lookup", agent_id=aid)
        row = get_agent_by_id(aid)
        if not row:
            _log("⚠️", "build_agents_metadata:agent_missing", agent_id=aid)
            continue
        training = load_training_data(aid)
        if len(training) > _TRAINING_SNIP_LEN:
            training = training[:_TRAINING_SNIP_LEN] + "\n...[truncated]"
        ds = ordered_available_datasources(row)
        api_entries = load_api_training_data(aid) if "API" in ds else []
        out.append(
            {
                "agent_id": str(aid),
                "agent_name": row.get("agent_name") or str(aid),
                "training_data_excerpt": training,
                "available_datasources": ds,
                "api_training_entries": api_entries,
            }
        )
        _log(
            "✅",
            "build_agents_metadata:agent_ready",
            agent_id=aid,
            datasources=ds,
            api_entries=len(api_entries),
        )
    _log("🏁", "build_agents_metadata:done", built=len(out))
    return out


def _format_agents_catalog_for_llm(agents_metadata: List[Dict[str, Any]]) -> str:
    blocks = []
    for m in agents_metadata:
        ds = ", ".join(m["available_datasources"])
        hint = m.get("recommended_datasource")
        if hint:
            hint_str = ", ".join(hint) if isinstance(hint, list) else str(hint)
            hint_line = f"Recommended datasource(s) for this question: {hint_str}\n"
            if isinstance(hint, list) and len(hint) > 1:
                hint_line += (
                    "NOTE: run this agent SEPARATELY for each datasource above - "
                    "output one sub_task per datasource.\n"
                )
        else:
            hint_line = ""
        blocks.append(
            f"---\n"
            f"Agent ID (use exactly): {m['agent_id']}\n"
            f"Agent name: {m['agent_name']}\n"
            f"Allowed datasources (pick one per sub-task): {ds}\n"
            f"{hint_line}"
            f"Training / schema / business context excerpt:\n{m['training_data_excerpt'] or '(no training text stored)'}\n"
        )
    return "\n".join(blocks)


def recommend_best_datasource_per_agent(
    question: str,
    agents_metadata: List[Dict[str, Any]],
    llm_agent: Any,
) -> Dict[str, List[str]]:
    """
    Single LLM call: decide best datasource per agent using typed evidence.
    - DATABASE evidence: initial training/schema text
    - API evidence: API training rows (name, endpoint, params, description)
    Returns {agent_id: ["DATABASE"]|["API"]|["FILE"]|["DATABASE","API"]}.
    Falls back to each agent's first allowed datasource on any error.
    """
    _log(
        "🔬",
        "recommend_best_datasource:start",
        agents=len(agents_metadata),
        question=question[:120],
    )

    fallback: Dict[str, List[str]] = {
        m["agent_id"]: (m["available_datasources"] or ["DATABASE"])[:1]
        for m in agents_metadata
    }
    if not agents_metadata:
        return fallback

    agent_ids_str = ", ".join(m["agent_id"] for m in agents_metadata)
    evidence_blocks: List[str] = []
    for m in agents_metadata:
        aid = str(m.get("agent_id") or "")
        aname = str(m.get("agent_name") or aid)
        allowed = m.get("available_datasources") or []
        db_evidence = str(m.get("training_data_excerpt") or "(no training text stored)")
        api_rows = m.get("api_training_entries") or []
        api_lines: List[str] = []
        for row in api_rows[:20]:
            if not isinstance(row, dict):
                continue
            api_lines.append(
                f"API: {str(row.get('name') or '').strip()}\n"
                f"  Endpoint: {str(row.get('endpoint') or '').strip()}\n"
                f"  Params: {str(row.get('base_params') or '').strip()}\n"
                f"  Description: {str(row.get('description') or '').strip()}"
            )
        api_evidence = "\n".join(api_lines) if api_lines else "(no API training rows)"
        evidence_blocks.append(
            f"---\n"
            f"Agent ID: {aid}\n"
            f"Agent name: {aname}\n"
            f"Available datasources: {', '.join(allowed)}\n\n"
            f"[DATABASE evidence]\n{db_evidence}\n\n"
            f"[API evidence]\n{api_evidence}\n"
        )
    evidence_catalog = "\n".join(evidence_blocks)

    sys_msg = (
        "You are a data-routing expert. For each agent, evaluate whether the question is best "
        "answered by DATABASE (schema/training evidence), API (endpoint/params evidence), or BOTH. "
        "If both datasources clearly and independently cover different aspects of the question, "
        "include both. Return only datasources from each agent's allowed list. "
        "Return JSON only, no markdown."
    )
    user_msg = (
        f"USER QUESTION:\n{question}\n\n"
        f"AGENT EVIDENCE:\n{evidence_catalog}\n\n"
        f"For each of these agent IDs ({agent_ids_str}), output the best datasource(s).\n"
        "Return JSON:\n"
        '{"recommendations": [{"agent_id": "<id>", "best_datasources": ["DATABASE|API|FILE"], '
        '"confidence": 0.0, "rationale": "brief"}]}'
    )

    try:
        _log_prompt_tokens("recommend_best_datasource", sys_msg, user_msg)
        raw = llm_agent.submit_prompt(
            [llm_agent.system_message(sys_msg), llm_agent.user_message(user_msg)],
            temperature=0.1,
            max_tokens=800,
        )
    except Exception as exc:
        _log("⚠️", "recommend_best_datasource:llm_error_using_fallback", error=str(exc))
        return fallback

    try:
        m = re.search(r"\{[\s\S]*\}", raw or "")
        data = json.loads(m.group(0)) if m else {}
        recs = data.get("recommendations") or []
        result = dict(fallback)
        allowed_map = {
            m["agent_id"]: set(m.get("available_datasources") or [])
            for m in agents_metadata
        }
        for r in recs:
            if not isinstance(r, dict):
                continue
            aid = str(r.get("agent_id") or "").strip()
            if aid not in allowed_map:
                _log(
                    "⚠️",
                    "recommend_best_datasource:invalid_or_unknown_agent",
                    agent_id=aid,
                    returned_ds=r.get("best_datasources") or r.get("best_datasource"),
                )
                continue

            raw_list = r.get("best_datasources") or []
            if not raw_list and r.get("best_datasource"):
                raw_list = [r.get("best_datasource")]

            valid_ds: List[str] = []
            for entry in raw_list:
                ds = str(entry or "").strip().upper()
                if ds == "DB":
                    ds = "DATABASE"
                if ds == "BOTH":
                    if "DATABASE" in allowed_map[aid]:
                        valid_ds.append("DATABASE")
                    if "API" in allowed_map[aid]:
                        valid_ds.append("API")
                    continue
                if ds in allowed_map[aid]:
                    valid_ds.append(ds)

            # keep order and remove duplicates
            deduped = list(dict.fromkeys(valid_ds))
            if deduped:
                result[aid] = deduped
                _log(
                    "💡",
                    "recommend_best_datasource:agent_recommendation",
                    agent_id=aid,
                    datasources=deduped,
                    confidence=r.get("confidence"),
                    rationale=str(r.get("rationale") or "")[:80],
                )
            else:
                _log(
                    "⚠️",
                    "recommend_best_datasource:invalid_or_unknown_agent",
                    agent_id=aid,
                    returned_ds=raw_list,
                )
        _log("✅", "recommend_best_datasource:done", recommendations=result)
        return result
    except Exception as exc:
        _log("⚠️", "recommend_best_datasource:parse_error_using_fallback", error=str(exc))
        return fallback


def _normalize_sub_tasks_from_llm_json(
    tasks: Any,
    allowed_ids: set,
) -> List[Dict[str, Any]]:
    if not isinstance(tasks, list):
        return []
    normalized: List[Dict[str, Any]] = []
    for t in tasks:
        if not isinstance(t, dict):
            continue
        sq = (t.get("sub_question") or t.get("question") or "").strip()
        aid = str(t.get("agent_id") or "").strip()
        ds = str(t.get("datasource") or "DATABASE").strip().upper()
        if ds == "DB":
            ds = "DATABASE"
        if not sq or aid not in allowed_ids:
            continue
        raw_score = t.get("confidence_score", None)
        if raw_score is None:
            raw_score = t.get("confidence", None)
        if raw_score is None:
            raw_score = 1.0
        try:
            score = float(raw_score)
        except (TypeError, ValueError):
            score = 1.0
        score = max(0.0, min(1.0, score))
        normalized.append(
            {
                "sub_question": sq,
                "agent_id": aid,
                "datasource": ds,
                "rationale": str(t.get("rationale") or ""),
                # Planner-provided confidence/relevance score (0.0–1.0).
                # Backward compatible: default to 1.0 if missing/invalid.
                "confidence_score": score,
            }
        )
    return normalized


def _parse_sub_tasks_json_response(raw_text: str, allowed_ids: set) -> List[Dict[str, Any]]:
    text = (raw_text or "").strip()
    m = re.search(r"\{[\s\S]*\}", text)
    if m:
        text = m.group(0)
    try:
        data = json.loads(text)
    except json.JSONDecodeError:
        print(f"[multi_agent_pipeline] sub_tasks JSON parse failed: {text[:500]}")
        return []
    tasks = data.get("sub_tasks") or data.get("subtasks") or []
    return _normalize_sub_tasks_from_llm_json(tasks, allowed_ids)


def decompose_question(
    question: str,
    agents_metadata: List[Dict[str, Any]],
    llm_agent: Any,
) -> List[Dict[str, Any]]:
    """
    LLM: split user question into sub_tasks using ONLY the provided agent catalog.
    Returns list of dicts with sub_question, agent_id, datasource, rationale.
    """
    _log("🧠", "decompose_question:start", question=question[:200], agents=len(agents_metadata))
    if not agents_metadata:
        _log("⚠️", "decompose_question:no_agents_metadata")
        return []

    allowed_ids = {m["agent_id"] for m in agents_metadata}
    catalog = _format_agents_catalog_for_llm(agents_metadata)

    selected_ids = ", ".join(sorted(allowed_ids))
    sys_msg = (
        "You are a senior data analyst planner. You MUST use only the agent IDs and "
        "datasource types listed in the catalog. Each sub-task maps to exactly one agent "
        "and one datasource. Split compound questions into independent sub-questions when "
        "different agents or domains apply. When one agent must answer using different "
        "systems (e.g. tabular data in DATABASE vs operational data in API vs documents in FILE), "
        "output multiple sub_tasks with the SAME agent_id but different datasource values. "
        "If the whole question is a single intent for one agent and one datasource, output exactly one sub_task. "
        "When MULTIPLE agents appear in the catalog, the user selected all of them deliberately: "
        "you MUST include at least one sub_task for EVERY agent whose training, name, or datasource "
        "capabilities can plausibly address part of the question."
        "mentions those domains, even if another agent could answer from SQL alone. "
        "Avoid over-broad sub_questions: they directly drive SQL/API generation. "
        "Respond with JSON only, no markdown. "
        "Don't include any agent name in the sub_question if any name is already present in the question then make sure not to remove it."
        "Each sub_question should be executable alone even without an agent name."

    )
    user_msg = (
        f"USER QUESTION:\n{question}\n\n"
        f"USER-SELECTED AGENT IDS (each must appear in sub_tasks when applicable): {selected_ids}\n\n"
        f"AGENT CATALOG:\n{catalog}\n\n"
        "Return JSON with this shape exactly:\n"
        '{"sub_tasks":[{"sub_question":"...","agent_id":"<id from catalog>",'
        '"datasource":"DATABASE|API|FILE",'
        '"confidence_score":0.0,"rationale":"brief"}]}\n'
        "Rules:\n"
        "- agent_id must match an Agent ID from the catalog character-for-character.\n"
        "- datasource must be allowed for that agent in the catalog.\n"
        "- confidence_score must be a float in the range 0.0–1.0 indicating how well this sub_question matches "
        "the agent's training/capabilities + user intent.\n"
        "- sub_question must be self-contained and executable alone.\n"
        "- You may repeat the same agent_id across sub_tasks when each needs a different datasource.\n"
        "- If several agents are listed, route intents to the best-matching agent; do not send "
        "API-oriented intents only to a generic DATABASE agent when another selected agent is clear.\n"
        "CRITICAL: Each sub_question must be MINIMAL and scoped to what the user literally asked. "
        "Do NOT add phrases like 'including all relevant customer/order details' or expand with extra "
        "columns/entities unless the user explicitly requested them. Narrow sub_questions produce "
        "narrower SQL/API results."
    )

    

    try:
        _log("📨", "decompose_question:llm_submit", allowed_ids=sorted(allowed_ids))
        _log_prompt_tokens("decompose_question", sys_msg, user_msg)
        raw = llm_agent.submit_prompt(
            [llm_agent.system_message(sys_msg), llm_agent.user_message(user_msg)],
            temperature=0.15,
            max_tokens=2500,
        )
    except Exception as e:
        _log("❌", "decompose_question:llm_error", error=str(e))
        return []

    text = (raw or "").strip()
    m = re.search(r"\{[\s\S]*\}", text)
    if m:
        text = m.group(0)
    try:
        data = json.loads(text)
    except json.JSONDecodeError:
        _log("❌", "decompose_question:json_parse_failed", preview=text[:500])
        return []

    tasks = data.get("sub_tasks") or data.get("subtasks") or []
    normalized = _normalize_sub_tasks_from_llm_json(tasks, allowed_ids)
    _log("✅", "decompose_question:done", raw_tasks=len(tasks) if isinstance(tasks, list) else 0, normalized=len(normalized))
    return normalized


def _llm_tasks_for_unused_agent(
    question: str,
    agent_meta: Dict[str, Any],
    tasks_assigned_elsewhere: List[Dict[str, Any]],
    llm_agent: Any,
) -> List[Dict[str, Any]]:
    """Second planning pass: assign intents to one agent that had zero sub_tasks."""
    aid = agent_meta["agent_id"]
    allowed_ids = {aid}
    catalog = _format_agents_catalog_for_llm([agent_meta])
    lines: List[str] = []
    for t in tasks_assigned_elsewhere:
        oaid = str(t.get("agent_id") or "")
        if oaid == aid:
            continue
        ds = t.get("datasource", "")
        sq = (t.get("sub_question") or "")[:400]
        lines.append(f"- Agent {oaid} [{ds}]: {sq}")
    block = "\n".join(lines) if lines else "(none)"

    sys_msg = (
        "You assign sub-tasks for ONE agent only. The user selected multiple agents; "
        "this agent currently has no work. Add sub_tasks that use ONLY this agent_id. "
        "Use its name and training: if it exposes APIs for trip, booking, shipment, or "
        "operational data, assign those question parts here rather than leaving them only "
        "to another agent. Do not copy tasks already fully covered by the same datasource "
        "on another agent unless this agent adds API access. JSON only, no markdown."
    )
    user_msg = (
        f"USER QUESTION:\n{question}\n\n"
        f"SUB-TASKS ALREADY ASSIGNED (other agents):\n{block}\n\n"
        f"PLAN ONLY FOR THIS AGENT (agent_id must be exactly {aid}):\n{catalog}\n\n"
        f'Return JSON: {{"sub_tasks":[{{"sub_question":"...","agent_id":"{aid}",'
        f'"datasource":"DATABASE|API|FILE","rationale":"brief"}}]}}\n'
        "If no part of the question fits this agent, return {\"sub_tasks\":[]}.\n"
    )
    try:
        _log_prompt_tokens("augment_agent_task", sys_msg, user_msg)
        raw = llm_agent.submit_prompt(
            [llm_agent.system_message(sys_msg), llm_agent.user_message(user_msg)],
            temperature=0.1,
            max_tokens=2000,
        )
    except Exception as e:
        print(f"[multi_agent_pipeline] _llm_tasks_for_unused_agent {aid}: {e}")
        return []
    return _parse_sub_tasks_json_response(raw, allowed_ids)


def _fallback_tasks_for_unused_agent(
    question: str,
    agent_meta: Dict[str, Any],
    llm_agent: Any,
) -> List[Dict[str, Any]]:
    """If the augment LLM returns nothing, use API intent or API-only agent fallback."""
    aid = agent_meta["agent_id"]
    avail: List[str] = agent_meta["available_datasources"]
    if not avail:
        return []
    only_api = avail == ["API"] or (len(avail) == 1 and avail[0] == "API")
    non_db = "DATABASE" not in avail and "API" in avail

    if only_api or non_db:
        return [
            {
                "sub_question": question,
                "agent_id": aid,
                "datasource": "API",
                "rationale": "fallback_unused_agent_api_preferred",
            }
        ]
    if "API" in avail:
        try:
            pick = llm_agent.select_datasource_from_intent(question, avail)
            ds = str(pick.get("datasource") or "").strip().upper()
            conf = float(pick.get("confidence") or 0.0)
            if ds == "API" and conf >= 0.5:
                return [
                    {
                        "sub_question": question,
                        "agent_id": aid,
                        "datasource": "API",
                        "rationale": f"fallback_intent_router:{pick.get('reasoning', '')}",
                    }
                ]
        except Exception as e:
            print(f"[multi_agent_pipeline] _fallback_tasks_for_unused_agent {aid}: {e}")
    return []


def augment_tasks_for_missing_agents(
    question: str,
    agents_metadata: List[Dict[str, Any]],
    existing: List[Dict[str, Any]],
    llm_agent: Any,
) -> List[Dict[str, Any]]:
    """
    Ensure every user-selected agent appears in the plan when multiple agents were chosen.
    Runs a focused LLM call per missing agent, then optional API/intent fallback.
    """
    _log("🧭", "augment_tasks_for_missing_agents:start", existing=len(existing), agents=len(agents_metadata))
    if len(agents_metadata) <= 1:
        _log("ℹ️", "augment_tasks_for_missing_agents:single_agent_skip")
        return existing
    required = {m["agent_id"] for m in agents_metadata}
    covered = {str(t.get("agent_id")) for t in existing if str(t.get("agent_id")) in required}
    missing_metas = [m for m in agents_metadata if m["agent_id"] not in covered]
    if not missing_metas:
        _log("✅", "augment_tasks_for_missing_agents:no_missing_agents")
        return existing

    out = list(existing)
    cumulative = list(existing)
    for meta in missing_metas:
        _log("➕", "augment_tasks_for_missing_agents:missing_agent", agent_id=meta["agent_id"])
        added = _llm_tasks_for_unused_agent(
            question, meta, cumulative, llm_agent
        )
        if not added:
            _log("🛟", "augment_tasks_for_missing_agents:llm_empty_try_fallback", agent_id=meta["agent_id"])
            added = _fallback_tasks_for_unused_agent(
                question, meta, llm_agent
            )
        for t in added:
            t["agent_id"] = meta["agent_id"]
        out.extend(added)
        cumulative = out
        _log("✅", "augment_tasks_for_missing_agents:agent_augmented", agent_id=meta["agent_id"], added=len(added))
    _log("🏁", "augment_tasks_for_missing_agents:done", total=len(out))
    return out


def build_multi_agent_datasource_map_from_tasks(
    sub_tasks: List[Dict[str, Any]],
) -> Dict[str, List[Dict[str, Any]]]:
    """Per-agent list of datasource assignments (multiple sub-tasks per agent supported)."""
    out: Dict[str, List[Dict[str, Any]]] = {}
    for t in sub_tasks:
        aid = t["agent_id"]
        out.setdefault(aid, []).append(
            {
                "datasource": t["datasource"],
                "sub_question": t["sub_question"],
                "rationale": t.get("rationale", ""),
            }
        )
    return out


def consolidate_same_agent_database_sub_tasks_for_mixed_plan(
    sub_tasks: List[Dict[str, Any]],
    original_question: str,
) -> List[Dict[str, Any]]:
    """
    When the plan mixes DATABASE with API/FILE, the decomposer often emits one DATABASE
    sub_task per metric (e.g. sales + booking) for the same agent, which runs SQL twice
    against the same schema. Collapse those into a single DATABASE sub_task per agent
    and use the *original* user question so one SQL can satisfy the combined intent.
    """
    if not sub_tasks or len(sub_tasks) < 2:
        return sub_tasks
    types = {str(t.get("datasource") or "").upper() for t in sub_tasks}
    if "DATABASE" not in types or not (types & {"API", "FILE"}):
        return sub_tasks
    by_agent: Dict[str, List[Dict[str, Any]]] = {}
    for t in sub_tasks:
        if str(t.get("datasource") or "").upper() == "DATABASE":
            aid = str(t.get("agent_id") or "")
            if not aid:
                continue
            by_agent.setdefault(aid, []).append(t)
    multi_db_agents = {aid for aid, tasks in by_agent.items() if len(tasks) > 1}
    if not multi_db_agents:
        return sub_tasks
    merged_first: Dict[str, Dict[str, Any]] = {}
    for aid in multi_db_agents:
        tasks = by_agent[aid]
        base = dict(tasks[0])
        oq = (original_question or "").strip()
        if oq:
            base["sub_question"] = oq
        else:
            base["sub_question"] = " ".join(
                (x.get("sub_question") or "").strip() for x in tasks
            )
        rats = [x.get("rationale") for x in tasks if x.get("rationale")]
        if rats:
            base["rationale"] = " | ".join(str(r) for r in rats)[:4000]
        scores: List[float] = []
        for x in tasks:
            cs = x.get("confidence_score", None)
            if cs is None:
                continue
            try:
                scores.append(float(cs))
            except (TypeError, ValueError):
                pass
        if scores:
            base["confidence_score"] = max(scores)
        merged_first[aid] = base
    out: List[Dict[str, Any]] = []
    seen_db_merged: set = set()
    for t in sub_tasks:
        aid = str(t.get("agent_id") or "")
        if str(t.get("datasource") or "").upper() == "DATABASE" and aid in multi_db_agents:
            if aid in seen_db_merged:
                continue
            seen_db_merged.add(aid)
            out.append(merged_first[aid])
        else:
            out.append(t)
    _log(
        "🔧",
        "consolidate_same_agent_database_sub_tasks_for_mixed_plan",
        before=len(sub_tasks),
        after=len(out),
        merged_agent_ids=sorted(multi_db_agents),
    )
    return out


def enhanced_datasource_resolver(
    sub_tasks: List[Dict[str, Any]],
    agents_metadata: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
    """
    Validate each sub-task: agent in roster; datasource must be allowed.
    If the model picked an invalid datasource, use that agent's first allowed datasource.
    """
    _log("🛣️", "enhanced_datasource_resolver:start", incoming=len(sub_tasks))
    by_id = {m["agent_id"]: m for m in agents_metadata}
    resolved: List[Dict[str, Any]] = []
    for t in sub_tasks:
        aid = t["agent_id"]
        meta = by_id.get(aid)
        if not meta:
            _log("⚠️", "enhanced_datasource_resolver:unknown_agent_skip", agent_id=aid)
            continue
        allowed: List[str] = meta["available_datasources"]
        requested_ds = t["datasource"]
        ds = requested_ds
        if ds not in allowed:
            ds = allowed[0]
            _log(
                "🔁",
                "enhanced_datasource_resolver:datasource_corrected",
                agent_id=aid,
                requested=requested_ds,
                selected=ds,
                allowed=allowed,
            )
        else:
            _log("✅", "enhanced_datasource_resolver:datasource_kept", agent_id=aid, datasource=ds)
        resolved.append(
            {
                "sub_question": t["sub_question"],
                "agent_id": aid,
                "datasource": ds,
                "rationale": t.get("rationale", ""),
            }
        )
    _log("🏁", "enhanced_datasource_resolver:done", resolved=len(resolved))
    return resolved


def filter_sub_tasks_by_threshold(
    sub_tasks: List[Dict[str, Any]],
    threshold: float,
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    """
    Split sub_tasks into (accepted, rejected) by confidence_score.
    Sub-tasks with missing/invalid confidence_score default to 1.0 upstream.
    """
    threshold = float(threshold)
    accepted: List[Dict[str, Any]] = []
    rejected: List[Dict[str, Any]] = []
    scores: List[float] = []
    for t in sub_tasks:
        raw_score = t.get("confidence_score", None)
        if raw_score is None:
            raw_score = 1.0
        score = float(raw_score)
        scores.append(score)
        if score >= threshold:
            accepted.append(t)
        else:
            rejected.append(t)

    if scores:
        mn = min(scores)
        mx = max(scores)
        avg = sum(scores) / len(scores)
    else:
        mn = mx = avg = 0.0

    _log(
        "🎚️",
        "filter_sub_tasks_by_threshold",
        threshold=threshold,
        accepted=len(accepted),
        rejected=len(rejected),
        min_score=round(mn, 4),
        max_score=round(mx, 4),
        avg_score=round(avg, 4),
    )

    # Log a bounded sample of rejected tasks so threshold issues are debuggable without spamming.
    if rejected:
        sample_n = 5
        rejected_sample = []
        for t in rejected[:sample_n]:
            rejected_sample.append(
                {
                    "agent_id": t.get("agent_id"),
                    "datasource": t.get("datasource"),
                    "confidence_score": t.get("confidence_score", 1.0),
                    "sub_question": (t.get("sub_question") or "")[:120],
                }
            )
        _log(
            "🧾",
            "filter_sub_tasks_by_threshold:rejected_sample",
            threshold=threshold,
            sample_count=min(sample_n, len(rejected)),
            rejected_sample=rejected_sample,
        )
    return accepted, rejected


def resolve_agent_and_datasource(
    sub_tasks: List[Dict[str, Any]],
    agents_metadata: List[Dict[str, Any]],
    datasource_override: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """
    Backward-compatible name. ``datasource_override`` is ignored; routing uses
    decomposition + training catalog only.
    """
    _ = datasource_override
    return enhanced_datasource_resolver(sub_tasks, agents_metadata)


def build_execution_plan(
    question: str,
    agent_ids: List[str],
    *,
    sql_to_use: Any = None,
    threshold: float = _DEFAULT_THRESHOLD,
    datasource_filter: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Training excerpts + LLM decomposition + validation.
    Returns keys: agents_metadata, raw_tasks, sub_tasks, multi_agent_datasource_map;
    or error dict with used_pipeline False.
    """
    from app.services.askme_service import get_askme_service

    _log("📐", "build_execution_plan:start", question=question[:200], agent_ids=agent_ids)
    agent_ids = [str(a) for a in agent_ids]
    agents_metadata = build_agents_metadata(agent_ids)
    if len(agents_metadata) < len(agent_ids):
        _log("❌", "build_execution_plan:missing_agent_metadata", requested=len(agent_ids), found=len(agents_metadata))
        return {
            "used_pipeline": False,
            "success": False,
            "reason": "missing_agent_metadata",
        }

    if datasource_filter:
        ds_norm = str(datasource_filter).strip().upper()
        if ds_norm == "DB":
            ds_norm = "DATABASE"
        before_count = len(agents_metadata)
        agents_metadata = [
            m for m in agents_metadata if ds_norm in (m.get("available_datasources") or [])
        ]
        _log(
            "🔽",
            "build_execution_plan:datasource_filter",
            requested=ds_norm,
            kept=len(agents_metadata),
            filtered_out=before_count - len(agents_metadata),
        )
        if not agents_metadata:
            _log(
                "⚠️",
                "build_execution_plan:no_agents_match_datasource_filter",
                datasource=ds_norm,
            )
            return {
                "used_pipeline": False,
                "success": False,
                "reason": "no_agents_match_datasource_filter",
            }

    svc0 = get_askme_service(agent_ids[0])
    if not svc0:
        _log("❌", "build_execution_plan:no_llm_service", agent_id=agent_ids[0])
        return {"used_pipeline": False, "success": False, "reason": "no_llm_service"}

    llm_agent = svc0._ensure_agent()
    ds_recommendations = recommend_best_datasource_per_agent(
        question, agents_metadata, llm_agent
    )
    for m in agents_metadata:
        m["recommended_datasource"] = ds_recommendations.get(m["agent_id"])
    _log(
        "🗺️",
        "build_execution_plan:datasource_recommendations",
        recommendations=ds_recommendations,
    )

    raw_tasks = decompose_question(question, agents_metadata, llm_agent)
    if raw_tasks:
        _log("🧩", "build_execution_plan:decomposed", count=len(raw_tasks))
        raw_tasks = augment_tasks_for_missing_agents(
            question, agents_metadata, raw_tasks, llm_agent
        )

    if not raw_tasks:
        _log("🛟", "build_execution_plan:fallback_single_task_no_decomposition")
        raw_tasks = [
            {
                "sub_question": question,
                "agent_id": agents_metadata[0]["agent_id"],
                "datasource": agents_metadata[0]["available_datasources"][0],
                "rationale": "fallback_single_task_no_llm_decomposition",
            }
        ]

    sub_tasks = enhanced_datasource_resolver(raw_tasks, agents_metadata)
    sub_tasks, rejected_tasks = filter_sub_tasks_by_threshold(sub_tasks, threshold)
    threshold_used = float(threshold)
    if rejected_tasks:
        _log(
            "🗑️",
            "build_execution_plan:tasks_filtered_out",
            threshold=threshold_used,
            rejected=len(rejected_tasks),
        )

    if not sub_tasks:
        _log("❌", "build_execution_plan:no_valid_sub_tasks_after_resolution")
        return {
            "used_pipeline": True,
            "success": False,
            "reason": "all_tasks_below_threshold",
            "sub_tasks": [],
            "raw_decomposition": raw_tasks,
            "agents_metadata": agents_metadata,
            "multi_agent_datasource_map": {},
            "rejected_tasks": rejected_tasks,
            "threshold_used": threshold_used,
        }

    sub_tasks = consolidate_same_agent_database_sub_tasks_for_mixed_plan(
        sub_tasks, question
    )
    multi_agent_datasource_map = build_multi_agent_datasource_map_from_tasks(sub_tasks)
    _log("✅", "build_execution_plan:done", raw_tasks=len(raw_tasks), sub_tasks=len(sub_tasks))
    pre_sql = (
        sql_to_use
        if len(sub_tasks) == 1 and sql_to_use
        and str(sql_to_use).strip().upper() not in ("N", "NONE", "[]")
        else None
    )
    return {
        "used_pipeline": True,
        "agents_metadata": agents_metadata,
        "raw_tasks": raw_tasks,
        "sub_tasks": sub_tasks,
        "multi_agent_datasource_map": multi_agent_datasource_map,
        "preloaded_sql_for_single_task": pre_sql,
        "rejected_tasks": rejected_tasks,
        "threshold_used": threshold_used,
    }


# Keywords that indicate a string is a real SQL query (not an LLM error message).
_SQL_START_KEYWORDS = (
    "SELECT", "WITH", "INSERT", "UPDATE", "DELETE", "MERGE", "CREATE", "DROP",
    "{",  # JSON query spec used by multi-query paths
)


def _is_valid_sql_candidate(value: Any) -> bool:
    """Return True if *value* looks like actual SQL or a JSON query spec.

    ``generate_sql_with_nodes`` returns human-readable error strings such as
    ``"SQL generation failed after N attempts..."`` when all retries are
    exhausted.  Passing those strings directly to the DB driver causes
    ORA-00900 (invalid SQL statement).  This helper detects the difference so
    callers can bail out cleanly instead of hitting the DB.
    """
    if not isinstance(value, str):
        return False
    stripped = value.strip()
    if not stripped:
        return False
    upper = stripped.upper()
    return any(upper.startswith(kw.upper()) for kw in _SQL_START_KEYWORDS)


def _normalize_sub_question_for_sql(sub_question: str) -> str:
    """Normalize punctuation copied from UI (en/em dash) so LLM/SQL gen sees ASCII."""
    if not sub_question:
        return sub_question
    return (
        sub_question.replace("\u2013", "-")
        .replace("\u2014", "-")
        .replace("\u2212", "-")
    )


def _unwrap_merge_bundle_to_agent_sql(generated: Any, agent_id: str) -> Any:
    """
    ``generate_sql_cached`` sometimes returns a merge payload (dict or JSON string) with
    ``agent_results`` instead of a single SQL string. Extract this agent's SQL so
    ``sql_trace["sql"]`` and DB execution use a plain query string (never the full merge JSON).
    """
    blob: Optional[Dict[str, Any]] = None
    list_blob: Optional[List[Dict[str, Any]]] = None
    if isinstance(generated, dict) and generated.get("agent_results"):
        blob = generated
    elif isinstance(generated, list):
        if generated and all(isinstance(x, dict) for x in generated):
            list_blob = generated
    elif isinstance(generated, str):
        s = generated.strip()
        if s.startswith(("{", "[")):
            try:
                d = json.loads(s)
            except json.JSONDecodeError:
                return generated
            if isinstance(d, dict) and d.get("agent_results"):
                blob = d
            elif isinstance(d, list) and d and all(isinstance(x, dict) for x in d):
                list_blob = d
    if blob is None:
        if list_blob is None:
            return generated
        blob = {"agent_results": list_blob}

    picked: Optional[str] = None
    for entry in blob.get("agent_results") or []:
        if not isinstance(entry, dict):
            continue
        if str(entry.get("agent_id")) == str(agent_id):
            sql_part = (entry.get("sql") or "").strip()
            err_part = entry.get("error")
            picked = sql_part if sql_part else (str(err_part) if err_part else None)
            break
    if not picked:
        for entry in blob.get("agent_results") or []:
            if isinstance(entry, dict) and (entry.get("sql") or "").strip():
                picked = str(entry.get("sql")).strip()
                break
    if picked:
        return picked
    return f"[agent {agent_id}] merge bundle contained no SQL text"


def execute_db_agent(
    agent_id: str, sub_question: str, preloaded_sql: Optional[Any] = None
) -> Tuple[Optional[pd.DataFrame], Any]:
    """Generate SQL for sub_question on this agent's DB and run it (single-agent service)."""
    from app.services.askme_service import get_askme_service

    sub_question = _normalize_sub_question_for_sql(sub_question or "")
    _log("🗄️", "execute_db_agent:start", agent_id=agent_id, question=sub_question[:180], has_preloaded_sql=bool(preloaded_sql))
    svc = get_askme_service(agent_id=agent_id)
    if not svc:
        _log("❌", "execute_db_agent:no_service", agent_id=agent_id)
        return None, f"[agent {agent_id}] AskmeService could not be loaded"
    try:
        svc._ensure_sql()
    except Exception as e:
        _log("❌", "execute_db_agent:ensure_sql_failed", agent_id=agent_id, error=str(e))
        return None, f"[agent {agent_id}] Database connection failed: {e}"

    if preloaded_sql:
        generated = preloaded_sql
    else:
        generated = svc.generate_sql_cached(question=sub_question)

    # Merge payload (dict or JSON string with agent_results) -> single SQL string for this agent.
    generated = _unwrap_merge_bundle_to_agent_sql(generated, agent_id)
    if isinstance(generated, str) and generated.startswith("[agent ") and "merge bundle" in generated:
        _log("⚠️", "execute_db_agent:no_sql_in_multi_bundle", agent_id=agent_id)
        return None, generated
    if isinstance(generated, dict) and generated.get("agent_results"):
        _log("⚠️", "execute_db_agent:unresolved_merge_dict", agent_id=agent_id)
        return None, f"[agent {agent_id}] Could not extract SQL from merge bundle"

    if not generated:
        _log("⚠️", "execute_db_agent:no_sql_generated", agent_id=agent_id)
        return None, f"[agent {agent_id}] SQL generation returned empty"

    # Guard: generate_sql_with_nodes returns a human-readable error string
    # (e.g. "SQL generation failed after N attempts...") when all retries are
    # exhausted.  That string is truthy but is not SQL — sending it to Oracle
    # causes ORA-00900.  Skip execution but **return the string** so sql_traces
    # and logs show why (previously we returned None and hid the failure).
    if isinstance(generated, str) and not _is_valid_sql_candidate(generated):
        _log("⚠️", "execute_db_agent:invalid_sql_candidate", agent_id=agent_id,
             preview=generated[:120])
        return None, generated

    if isinstance(generated, str):
        try:
            parsed = json.loads(generated)
            if isinstance(parsed, dict) and "queries" in parsed:
                df = svc.run_sql_cached(
                    generated, merge_info=parsed, question=sub_question
                )
            else:
                df = svc.run_sql_cached(generated, question=sub_question)
        except json.JSONDecodeError:
            df = svc.run_sql_cached(generated, question=sub_question)
    else:
        df = svc.run_sql_cached(str(generated), question=sub_question)

    if isinstance(df, dict):
        parts = [
            x for x in df.values() if x is not None and not getattr(x, "empty", True)
        ]
        if not parts:
            return None, generated
        df = pd.concat(parts, axis=0, ignore_index=True)

    if df is None or getattr(df, "empty", True):
        _log("⚠️", "execute_db_agent:empty_dataframe", agent_id=agent_id)
        return None, generated
    _log("✅", "execute_db_agent:done", agent_id=agent_id, rows=int(len(df.index)))
    return df, generated


def execute_api_agent(
    agent_id: str, sub_question: str,
) -> Tuple[Optional[pd.DataFrame], Dict[str, Any]]:
    """Run API path for one agent (datasource auto-selected inside process_api_user_query)."""
    from app.services.askme_service import get_askme_service

    _log("🌐", "execute_api_agent:start", agent_id=agent_id, question=sub_question[:180])
    svc = get_askme_service(agent_id=agent_id)
    if not svc:
        _log("❌", "execute_api_agent:no_service", agent_id=agent_id)
        return None, {"success": False, "message": "no service"}
    session_id = f"multi_agent_pipeline_{agent_id}"
    try:
        res = svc.process_api_user_query(
            sub_question, session_id=session_id, agent_id=agent_id
        )
    except Exception as e:
        _log("❌", "execute_api_agent:exception", agent_id=agent_id, error=str(e))
        return None, {"success": False, "message": str(e)}

    if not res or not res.get("success"):
        _log("⚠️", "execute_api_agent:unsuccessful_response", agent_id=agent_id, success=bool(res and res.get("success")))
        return None, res or {}

    df = res.get("df")
    if isinstance(df, list) and df:
        frames = [x for x in df if isinstance(x, pd.DataFrame) and not x.empty]
        if frames:
            df = pd.concat(frames, axis=0, ignore_index=True)
        else:
            df = None
    if df is None or (hasattr(df, "empty") and df.empty):
        _log("⚠️", "execute_api_agent:empty_dataframe", agent_id=agent_id)
        return None, res
    _log("✅", "execute_api_agent:done", agent_id=agent_id, rows=int(len(df.index)))
    return df, res


def _slim_api_process_result_for_trace(api_res: Any) -> Dict[str, Any]:
    """
    process_api_user_query returns raw_data, df, long message, etc. Those must not be
    embedded in api_traces (JSON responses become tens of MB). Match askme
    _api_result_dict_for_envelope: api_endpoint, api_params, api_response_validator only.
    """
    if not isinstance(api_res, dict):
        return {"api_endpoint": "", "api_params": None, "api_response_validator": ""}
    ep = api_res.get("api_endpoint")
    val = api_res.get("code")
    pr = api_res.get("params")
    base: Dict[str, Any] = {
        "api_endpoint": ep if ep is not None else "",
        "api_params": pr,
        "api_response_validator": "" if val is None else val,
    }
    if not api_res.get("success"):
        base["success"] = False
        msg = api_res.get("message") or ""
        base["message"] = msg[:4000] if len(str(msg)) > 4000 else msg
    return base


def normalize_to_dataframe(
    obj: Any, label: str = "value"
) -> Optional[pd.DataFrame]:
    """Best-effort: API/DB outputs → DataFrame with simple column normalization."""
    if obj is None:
        return None
    if isinstance(obj, pd.DataFrame):
        if obj.empty:
            return None
        out = obj.copy()
        out.columns = [str(c).strip().replace(" ", "_") for c in out.columns]
        return out
    if isinstance(obj, list) and obj and isinstance(obj[0], dict):
        return pd.DataFrame(obj)
    if isinstance(obj, dict):
        return pd.DataFrame([{label: json.dumps(obj, default=str)}])
    return pd.DataFrame([{label: str(obj)}])


def _append_normalized_tabular(
    tabular_frames: List[pd.DataFrame], df: Optional[pd.DataFrame]
) -> bool:
    """
    Append a tabular frame without using ``normalize(df) or df`` — pandas forbids
    ``bool(DataFrame)``, which Python applies to the left operand of ``or``.
    """
    if df is None:
        return False
    ndf = normalize_to_dataframe(df)
    if ndf is not None:
        tabular_frames.append(ndf)
        return True
    if isinstance(df, pd.DataFrame) and not df.empty:
        tabular_frames.append(df)


# ---------------------------------------------------------------------------
# Merge strategy inference
# ---------------------------------------------------------------------------

_STRATEGY_KEYWORDS: Dict[str, List[str]] = {
    "forecast": ["forecast", "predict", "projection", "trend", "future"],
    "compare":  ["compare", "comparison", "vs", "versus", "difference", "diff", "contrast"],
    "union":    ["combine", "union", "together", "all", "combined", "aggregate", "list all"],
    "join":     ["join", "related", "link", "match", "common", "shared", "merge"],
}


def infer_merge_strategy(question: str, frames: List[pd.DataFrame]) -> str:
    """
    Infer the appropriate merge strategy from the question text and frame schemas.

    Priority order: explicit keyword match → schema analysis → "union" fallback.

    Returns one of: "single" | "union" | "join" | "compare" | "forecast"
    """
    if len(frames) <= 1:
        return "single"

    q_lower = question.lower()
    for strategy, keywords in _STRATEGY_KEYWORDS.items():
        if any(kw in q_lower for kw in keywords):
            return strategy

    # Schema-based heuristic: if frames share >30 % of columns try join, else union
    cols = [set(str(c).lower() for c in df.columns) for df in frames if df is not None]
    if len(cols) >= 2:
        common = cols[0].intersection(cols[1])
        union_size = cols[0].union(cols[1])
        if union_size and len(common) / len(union_size) > 0.30:
            return "join"

    return "union"


# ---------------------------------------------------------------------------
# Core merge engine
# ---------------------------------------------------------------------------

def _infer_join_key(left: pd.DataFrame, right: pd.DataFrame) -> Optional[str]:
    """Return the best single join column (case-insensitive match, prefer ID-like names)."""
    lc = {str(c).lower(): c for c in left.columns}
    rc = {str(c).lower(): c for c in right.columns}
    common_lower = [k for k in lc if k in rc]
    if not common_lower:
        return None
    id_like = [k for k in common_lower if any(
        hint in k for hint in ("id", "key", "code", "no", "number", "ref")
    )]
    best = id_like[0] if id_like else common_lower[0]
    return lc[best]  # return original column name from left


def _normalize_columns(df: pd.DataFrame) -> pd.DataFrame:
    """Standardise column names: strip, lowercase, replace spaces with underscores."""
    df = df.copy()
    df.columns = [str(c).strip().lower().replace(" ", "_") for c in df.columns]
    return df


def merge_dataframes(
    frames: List[pd.DataFrame],
    strategy: str = "auto",
    question: str = "",
    max_rows: int = 5000,
) -> Tuple[pd.DataFrame, Dict[str, Any]]:
    """
    Merge a list of DataFrames according to *strategy*.

    Parameters
    ----------
    frames:    Non-empty list of DataFrames (None / empty entries are skipped).
    strategy:  One of "single" | "union" | "join" | "compare" | "forecast" | "auto".
               "auto" → infer from *question* and schema.
    question:  Original user question (used only when strategy == "auto").
    max_rows:  Hard cap on returned rows.

    Returns
    -------
    (merged_df, merge_meta)
    merge_meta includes: merge_strategy, strategies (list of step dicts), truncated.
    """
    valid = [df for df in frames if df is not None and not getattr(df, "empty", True)]
    meta: Dict[str, Any] = {"strategies": [], "truncated": False}

    if not valid:
        meta["merge_strategy"] = "empty"
        return pd.DataFrame(), meta

    if len(valid) == 1 or strategy == "single":
        meta["merge_strategy"] = "single"
        result = valid[0].head(max_rows)
        meta["truncated"] = len(valid[0]) > max_rows
        return result, meta

    # Resolve "auto"
    if strategy == "auto":
        strategy = infer_merge_strategy(question, valid)

    meta["merge_strategy"] = strategy

    # Normalise column names before any merge to handle schema heterogeneity
    norm = [_normalize_columns(df) for df in valid]

    try:
        if strategy == "union":
            merged = pd.concat(norm, axis=0, ignore_index=True)
            meta["strategies"].append({"step": "concat_rows", "frames": len(norm)})

        elif strategy == "join":
            merged = norm[0]
            for nxt in norm[1:]:
                key = _infer_join_key(merged, nxt)
                if key:
                    merged = pd.merge(merged, nxt, on=key, how="outer", suffixes=("", "_r"))
                    meta["strategies"].append({"step": "merge_on_key", "key": key})
                else:
                    # No common key → fall back to column-wise concat (side-by-side)
                    merged = pd.concat(
                        [merged.reset_index(drop=True), nxt.reset_index(drop=True)], axis=1
                    )
                    meta["strategies"].append({"step": "concat_columns_no_key"})

        elif strategy == "compare":
            # Label each frame with its source index, then union for side-by-side comparison
            labelled = []
            for i, df in enumerate(norm):
                df = df.copy()
                df["_source"] = f"agent_{i + 1}"
                labelled.append(df)
            merged = pd.concat(labelled, axis=0, ignore_index=True)
            meta["strategies"].append({"step": "union_with_source_label", "frames": len(norm)})

        elif strategy == "forecast":
            # Forecast operates on the primary frame; extras are ignored or appended
            merged = pd.concat(norm, axis=0, ignore_index=True)
            meta["strategies"].append({"step": "concat_rows_for_forecast", "frames": len(norm)})

        else:
            # Unknown strategy — safe fallback: union
            merged = pd.concat(norm, axis=0, ignore_index=True)
            meta["strategies"].append({"step": "fallback_union", "unknown_strategy": strategy})
            meta["merge_strategy"] = "union"

    except Exception as exc:
        _log("❌", "merge_dataframes:error", strategy=strategy, error=str(exc))
        # Last-resort fallback: return first valid frame unchanged
        meta["strategies"].append({"step": "error_fallback", "error": str(exc)})
        meta["merge_strategy"] = "error_fallback"
        merged = valid[0]

    truncated = len(merged) > max_rows
    if truncated:
        merged = merged.head(max_rows)
    meta["truncated"] = truncated
    return merged, meta


# ---------------------------------------------------------------------------
# Response builder
# ---------------------------------------------------------------------------

def build_response(
    merged_df: Optional[pd.DataFrame],
    merge_meta: Dict[str, Any],
    *,
    sub_tasks: Optional[List[Dict[str, Any]]] = None,
    file_results: Optional[List[Dict[str, Any]]] = None,
    question: str = "",
    api_traces: Optional[List[Dict[str, Any]]] = None,
) -> Dict[str, Any]:
    """
    Convert a merged DataFrame into the standard response payload.

    Always returns:
        df.headers  – list of human-readable column names
        df.data     – list of row dicts
        merge_meta  – merge diagnostics
        file_results, sub_tasks, api_traces – passed through unchanged

    Backward compatible: callers that already build response_details manually
    can continue to do so; this helper is additive.
    """
    import math

    def _clean(val: Any) -> Any:
        """Make a scalar JSON-safe."""
        if isinstance(val, float) and (math.isnan(val) or math.isinf(val)):
            return None
        return val

    def _clean_records(records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        return [{k: _clean(v) for k, v in row.items()} for row in records]

    has_tabular = (
        merged_df is not None
        and isinstance(merged_df, pd.DataFrame)
        and not merged_df.empty
    )

    if has_tabular:
        raw_columns = list(merged_df.columns)
        headers = [str(c).replace("_", " ").capitalize() for c in raw_columns]
        data = _clean_records(merged_df.to_dict("records"))
    else:
        headers = []
        data = []

    merge_strategy = merge_meta.get("merge_strategy", "unknown")
    agent_count = len(sub_tasks) if sub_tasks else 0

    return {
        "df": {
            "headers": headers,
            "data": data,
            "merge_strategy": merge_strategy,
            "multi_agent_mixed": agent_count > 1,
            "tabular_from_merge": has_tabular,
            "training_driven_pipeline": True,
        },
        "merge_meta": merge_meta,
        "file_results": file_results or [],
        "sub_tasks": sub_tasks or [],
        "api_traces": api_traces or [],
        "question": question,
    }


# ---------------------------------------------------------------------------
# merge_results — original heterogeneous behaviour (backward compatible).
# Call merge_dataframes() directly when an explicit strategy is needed.
# ---------------------------------------------------------------------------

def merge_results(
    dataframes: List[pd.DataFrame],
    max_rows: int = 5000,
    *,
    strategy: str = "auto",
    question: str = "",
    frame_agent_ids: Optional[List[str]] = None,
) -> Tuple[pd.DataFrame, Dict[str, Any]]:
    """
    Merge ordered tabular results (DB/API only).

    When called without explicit *strategy* / *question* (the default from
    execute_plan_async), the function preserves the original
    heterogeneous_merge_dataframes behaviour so existing callers are not
    affected.  Pass an explicit *strategy* to use the new merge engine.
    """
    from app.services.askme_service import heterogeneous_merge_dataframes

    agents = list(frame_agent_ids or [])
    filtered_pairs: List[Tuple[pd.DataFrame, str]] = []
    for idx, d in enumerate(dataframes):
        if d is None or getattr(d, "empty", True):
            continue
        aid = agents[idx] if idx < len(agents) else ""
        filtered_pairs.append((d, str(aid)))
    dfs = [d for d, _ in filtered_pairs]
    filtered_agent_ids = [aid for _, aid in filtered_pairs]
    if not dfs:
        return pd.DataFrame(), {"merge_strategy": "empty", "strategies": []}
    if len(dfs) == 1:
        return dfs[0].head(max_rows), {"merge_strategy": "single", "strategies": []}
    if _all_frames_disjoint(dfs):
        _log("🔀", "merge_results:disjoint_keep_separate", count=len(dfs))
        return pd.DataFrame(), {
            "merge_strategy": "separate_frames",
            "strategies": [],
            "separate_frames": [
                {"agent_id": filtered_agent_ids[i], "df": dfs[i].head(max_rows)}
                for i in range(len(dfs))
            ],
        }
    # Explicit strategy requested → use new merge engine
    if strategy != "auto" or question:
        return merge_dataframes(dfs, strategy=strategy, question=question, max_rows=max_rows)

    # Default path: original heterogeneous pairwise merge (preserves column names/types)
    merged = dfs[0]
    strategies: List[Dict[str, Any]] = []
    for nxt in dfs[1:]:
        merged, meta = heterogeneous_merge_dataframes(merged, nxt, max_rows=max_rows)
        strategies.append(meta)
    return merged, {"merge_strategy": "heterogeneous_chain", "strategies": strategies}


def combine_agent_results(
    dataframes: List[pd.DataFrame],
    max_rows: int = 5000,
) -> Tuple[pd.DataFrame, Dict[str, Any]]:
    """Alias for merge_results (multi-table agent outputs)."""
    return merge_results(dataframes, max_rows=max_rows)


def execute_multi_datasource_single_agent(
    agent_id: str,
    sub_tasks: List[Dict[str, Any]],
    *,
    preloaded_sql: Optional[Any] = None,
    max_rows: int = 5000,
) -> Dict[str, Any]:
    """Run only sub_tasks for one agent; merges tabular slices like full execute_plan."""
    filtered = [t for t in sub_tasks if t["agent_id"] == agent_id]
    return execute_plan(
        filtered,
        preloaded_sql=preloaded_sql,
        max_rows=max_rows,
    )


def score_execution_result(
    df: Optional[pd.DataFrame],
    datasource: str,
) -> float:
    """
    Heuristic result-quality score based on rows returned.

    This is intentionally lightweight (no extra LLM calls).
    """
    if str(datasource).upper() == "FILE":
        # FILE outputs are payloads, not tabular results to be merged.
        return 1.0
    if df is None or getattr(df, "empty", True):
        return 0.0
    n = int(len(df.index))
    if n >= 100:
        return 1.0
    if n >= 10:
        return 0.8
    return 0.6


def _execute_single_sub_task_sync(
    index: int,
    st: Dict[str, Any],
    *,
    preloaded_sql: Optional[Any] = None,
    context_agent_ids: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """
    Execute one sub-task synchronously and return a normalized record that can be
    deterministically re-ordered after concurrent execution.
    """
    from app.services.askme_service import run_file_agent_document_query_sync

    aid = st["agent_id"]
    ds = st["datasource"]
    sq = st["sub_question"]
    out: Dict[str, Any] = {
        "index": index,
        "tabular_df": None,
        "sql_trace": None,
        "api_trace": None,
        "file_result": None,
        "result_score": 0.0,
    }

    _log("🚦", "sub_task:start", index=index, agent_id=aid, datasource=ds, question=sq[:180])
    if ds == "DATABASE":
        _log("➡️", "sub_task:route_database", index=index, agent_id=aid)
        df, sql_g = execute_db_agent(aid, sq, preloaded_sql=preloaded_sql)
        out["sql_trace"] = {"agent_id": aid, "sub_question": sq, "sql": sql_g}
        if df is not None:
            out["tabular_df"] = df
        out["result_score"] = score_execution_result(out["tabular_df"], ds)
        return out

    if ds == "API":
        _log("➡️", "sub_task:route_api", index=index, agent_id=aid)
        df, api_res = execute_api_agent(aid, sq)
        out["api_trace"] = {
            "agent_id": aid,
            "sub_question": sq,
            "result": _slim_api_process_result_for_trace(api_res),
        }
        if df is not None:
            out["tabular_df"] = df
        out["result_score"] = score_execution_result(out["tabular_df"], ds)
        return out

    if ds == "FILE":
        _log("➡️", "sub_task:route_file", index=index, agent_id=aid)
        try:
            # For FILE datasources, we broaden the search to all context agent IDs 
            # if available, ensuring cross-agent retrieval when 'all agents' are selected.
            target_agent_id = context_agent_ids if context_agent_ids else aid
            fr = run_file_agent_document_query_sync(target_agent_id, sq)
            if fr:
                out["file_result"] = fr
        except Exception as e:
            out["file_result"] = {
                "agent_id": aid,
                "datasource": "FILE",
                "document_result_type": "error",
                "document_error": {"error": str(e)},
                "question": sq,
                "document_query": "",
                "sql": "N",
                "api_endpoint": "",
                "api_params": "",
                "api_response_validator": "",
            }
        out["result_score"] = score_execution_result(None, ds)
        return out

    _log("⚠️", "sub_task:unknown_datasource_skip", index=index, agent_id=aid, datasource=ds)
    return out


async def _execute_plan_async(
    sub_tasks: List[Dict[str, Any]],
    *,
    preloaded_sql: Optional[Any] = None,
    max_rows: int = 5000,
    context_agent_ids: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """
    Execute validated sub_tasks concurrently; preserve original ordering in traces
    and frame collection before merge.
    """
    _log("⚙️", "execute_plan_async:start", sub_tasks=len(sub_tasks), max_rows=max_rows)
    tabular_frames: List[pd.DataFrame] = []
    tabular_agent_ids: List[str] = []
    file_results: List[Dict[str, Any]] = []
    sql_traces: List[Any] = []
    api_traces: List[Dict[str, Any]] = []

    only_pre_sql = preloaded_sql if len(sub_tasks) == 1 else None
    tasks = [
        asyncio.to_thread(
            _execute_single_sub_task_sync,
            idx,
            st,
            preloaded_sql=only_pre_sql if idx == 0 else None,
            context_agent_ids=context_agent_ids,
        )
        for idx, st in enumerate(sub_tasks)
    ]
    _log("🧵", "execute_plan_async:tasks_dispatched", count=len(tasks))
    results = await asyncio.gather(*tasks)
    results.sort(key=lambda x: x["index"])
    _log("📥", "execute_plan_async:results_collected", count=len(results))

    for row in results:
        df = row.get("tabular_df")
        score = float(row.get("result_score") or 0.0)
        aid = str((row.get("sql_trace") or row.get("api_trace") or {}).get("agent_id") or "")
        if df is not None and score >= _RESULT_SCORE_THRESHOLD:
            appended = _append_normalized_tabular(tabular_frames, df)
            if appended:
                tabular_agent_ids.append(aid)
        elif df is not None:
            _log(
                "🗑️",
                "execute_plan_async:result_filtered_low_score",
                agent_id=aid,
                score=score,
                threshold=_RESULT_SCORE_THRESHOLD,
            )

        sql_t = row.get("sql_trace")
        if sql_t is not None:
            sql_traces.append(sql_t)

        api_t = row.get("api_trace")
        if api_t is not None:
            api_traces.append(api_t)

        fr = row.get("file_result")
        if fr is not None:
            file_results.append(fr)

    merged, merge_meta = merge_results(
        tabular_frames,
        max_rows=max_rows,
        frame_agent_ids=tabular_agent_ids,
    )
    separate_frames = merge_meta.get("separate_frames") or []

    _ok_file = any(
        fr.get("document_result_type")
        in ("document_search", "document_analysis", "document_find")
        for fr in file_results
    )
    has_tabular = (
        merged is not None and not getattr(merged, "empty", True)
    ) or bool(separate_frames)
    success = has_tabular or _ok_file
    _log(
        "🏁",
        "execute_plan_async:done",
        success=success,
        tabular_frames=len(tabular_frames),
        sql_traces=len(sql_traces),
        api_traces=len(api_traces),
        file_results=len(file_results),
    )

    return {
        "success": success,
        "df": merged,
        "file_results": file_results,
        "merge_meta": merge_meta,
        "separate_dfs": separate_frames,
        "sql_traces": sql_traces,
        "api_traces": api_traces,
    }


def execute_plan(
    sub_tasks: List[Dict[str, Any]],
    *,
    preloaded_sql: Optional[Any] = None,
    max_rows: int = 5000,
    context_agent_ids: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """
    Execute validated sub_tasks in order; collect SQL/API traces and optional file payloads.
    """
    _log("🚀", "execute_plan:start", sub_tasks=len(sub_tasks), max_rows=max_rows)
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        _log("🟢", "execute_plan:no_running_loop_using_asyncio_run")
        return asyncio.run(
            _execute_plan_async(
                sub_tasks,
                preloaded_sql=preloaded_sql,
                max_rows=max_rows,
                context_agent_ids=context_agent_ids,
            )
        )

    # Keep sync API usable even if called from a thread with an active event loop.
    _log("🟡", "execute_plan:running_loop_detected_using_threadpool")
    with ThreadPoolExecutor(max_workers=1) as pool:
        fut = pool.submit(
            asyncio.run,
            _execute_plan_async(
                sub_tasks,
                preloaded_sql=preloaded_sql,
                max_rows=max_rows,
                context_agent_ids=context_agent_ids,
            ),
        )
        return fut.result()


def handle_mixed_datasource_execution(
    sub_question_map: List[Dict[str, Any]],
    *,
    preloaded_sql: Optional[Any] = None,
    max_rows: int = 5000,
) -> Dict[str, Any]:
    """
    Entry point for mixed DB/API/FILE sub-task lists (homogeneous DB-only or API-only
    still flows through this function; route-level code may short-circuit to legacy paths).
    """
    return execute_plan(
        sub_question_map,
        preloaded_sql=preloaded_sql,
        max_rows=max_rows,
    )


def multi_agent_pipeline(
    question: str,
    agent_ids: List[str],
    *,
    sql_to_use: Any = None,
    max_rows: int = 5000,
    threshold: float = _DEFAULT_THRESHOLD,
    datasource_filter: Optional[str] = None,
) -> Dict[str, Any]:
    """
    End-to-end training-driven pipeline for multiple agents.

    Returns:
        used_pipeline: True if planner ran and produced validated sub_tasks (even one).
        success: True if any tabular or file result succeeded.
        df: merged tabular DataFrame (may be empty).
        file_results: list of file agent payloads (same as ask-file style).
        sub_tasks: validated decomposition for explainability.
        multi_agent_datasource_map: agent_id -> list of {datasource, sub_question, rationale}.
        sql_traces / api_traces: per sub-task execution metadata.
    """
    _log(
        "🎯",
        "multi_agent_pipeline:start",
        question=question[:200],
        agent_ids=agent_ids,
        datasource_filter=datasource_filter,
    )
    plan = build_execution_plan(
        question,
        agent_ids,
        sql_to_use=sql_to_use,
        threshold=threshold,
        datasource_filter=datasource_filter,
    )
    if not plan.get("used_pipeline"):
        _log("❌", "multi_agent_pipeline:planner_not_used", reason=plan.get("reason"))
        return plan
    if plan.get("reason") in ("no_valid_sub_tasks_after_resolution", "all_tasks_below_threshold"):
        return {
            "used_pipeline": True,
            "success": False,
            "reason": plan.get("reason"),
            "sub_tasks": plan.get("sub_tasks", []),
            "raw_decomposition": plan.get("raw_decomposition", []),
            "multi_agent_datasource_map": plan.get("multi_agent_datasource_map", {}),
            "rejected_tasks": plan.get("rejected_tasks", []),
            "threshold_used": plan.get("threshold_used", threshold),
        }

    sub_tasks = plan["sub_tasks"]
    multi_agent_datasource_map = plan["multi_agent_datasource_map"]
    raw_tasks = plan["raw_tasks"]
    exec_out = execute_plan(
        sub_tasks,
        preloaded_sql=plan.get("preloaded_sql_for_single_task"),
        max_rows=max_rows,
        context_agent_ids=agent_ids,
    )

    merged = exec_out["df"]
    merge_meta = exec_out["merge_meta"]
    separate_dfs = exec_out.get("separate_dfs") or []
    file_results = exec_out["file_results"]
    sql_traces = exec_out["sql_traces"]
    api_traces = exec_out["api_traces"]
    success = exec_out["success"]

    # ------------------------------------------------------------------
    # 🪙 Pipeline-level token budget summary
    # ------------------------------------------------------------------
    # Count the planner prompts that were built during build_execution_plan.
    # We re-derive them here (strings are already constructed above) so we
    # can emit a clean summary without threading a tracker through every call.
    _plan_llm_calls = 1 + max(0, len(sub_tasks) - 1)  # decompose + augment passes
    _question_tokens = _count_tokens(question)
    _log(
        "💰",
        "token_budget:pipeline_summary",
        question_tokens=_question_tokens,
        estimated_planner_llm_calls=_plan_llm_calls,
        sub_tasks_executed=len(sub_tasks),
        tiktoken_backend="tiktoken" if _TIKTOKEN_AVAILABLE else "word-estimate",
    )

    _log(
        "🏆",
        "multi_agent_pipeline:done",
        success=success,
        sub_tasks=len(sub_tasks),
        has_df=bool(merged is not None and not getattr(merged, "empty", True)),
        file_results=len(file_results),
    )

    dm_list = [
        {
            "agent_id": t["agent_id"],
            "type": t["datasource"],
            "confidence": t.get("confidence_score", 1.0),
            "reasoning": t.get("rationale", "training_driven_pipeline"),
            "sub_question": t["sub_question"],
        }
        for t in sub_tasks
    ]
    return {
        "used_pipeline": True,
        "success": success,
        "df": merged,
        "separate_dfs": separate_dfs,
        "file_results": file_results,
        "sub_tasks": sub_tasks,
        "raw_decomposition": raw_tasks,
        "merge_meta": merge_meta,
        "sql_traces": sql_traces,
        "api_traces": api_traces,
        "datasource_map_list": dm_list,
        "multi_agent_datasource_map": multi_agent_datasource_map,
        "rejected_tasks": plan.get("rejected_tasks", []),
        "threshold_used": plan.get("threshold_used", threshold),
    }
