from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

from app.services.talkback.utils import (
    _collect_api_names,
    _df_row_count_and_columns,
    _extract_talkback_options,
    _infer_datasource_from_details,
    _parse_sql_preview_for_meta,
)

logger = logging.getLogger(__name__)


def _dispatch_agents_for_talkback_meta(agent_ids: Optional[List[str]]) -> Optional[List[Dict[str, str]]]:
    """Stable id + display name for history (avoid encoding agent names in ``question``)."""
    if not agent_ids:
        return None
    from app.services.agent_config import get_agent_by_id

    out: List[Dict[str, str]] = []
    for aid in agent_ids[:12]:
        s = str(aid).strip()
        if not s:
            continue
        row = get_agent_by_id(s)
        name = (row.get("agent_name") if row else None) or s
        out.append({"agent_id": s, "agent_name": str(name)})
    return out or None


def _agent_display_labels(agent_ids: Optional[List[str]]) -> List[str]:
    """Resolve agent ids to display names (fallback to id)."""
    if not agent_ids:
        return []
    from app.services.agent_config import get_agent_by_id

    labels: List[str] = []
    for aid in agent_ids:
        s = str(aid).strip()
        if not s:
            continue
        row = get_agent_by_id(s)
        labels.append(str((row.get("agent_name") if row else None) or s))
    return labels


def build_talkback_answer_meta(
    result: Dict[str, Any],
    *,
    dispatch_agent_ids: Optional[List[str]] = None,
    clarification_selection: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Compact JSON-serializable metadata from an ask response for talkback history.
    Improves context relevance and routing rewrite by preserving intent, datasource, and
    data shape without storing full dataframes.
    """
    out: Dict[str, Any] = {}
    if isinstance(result.get("success"), bool):
        out["success"] = result["success"]
    if result.get("direct_prior_df_llm_bypass"):
        out["prior_df_bypass"] = True

    details = result.get("details")
    if not isinstance(details, dict):
        details = {}

    # Extract suggested questions from talkback string if present
    tb_text = details.get("talkback")
    if tb_text:
        opts = _extract_talkback_options(str(tb_text))
        if opts:
            out["talkback_options"] = opts

    ir = result.get("intent_routing")
    if isinstance(ir, dict):
        if ir.get("intent") is not None:
            out["intent"] = str(ir.get("intent"))[:120]
        sug = ir.get("suggested_agent_ids")
        if isinstance(sug, list) and sug:
            out["suggested_agent_ids"] = [str(x) for x in sug[:8]]
        if bool(ir.get("needs_clarification")):
            out["needs_clarification"] = True
        clar_ids = ir.get("clarification_agent_ids")
        if isinstance(clar_ids, list) and clar_ids:
            out["clarification_agent_ids"] = [str(x) for x in clar_ids[:12]]

    ds = _infer_datasource_from_details(details)
    if ds:
        out["datasource"] = ds

    sp = _parse_sql_preview_for_meta(details.get("sql"))
    if sp:
        out["sql_preview"] = sp

    n, cols = _df_row_count_and_columns(details)
    if n is not None:
        out["row_count"] = n
    if cols:
        out["columns"] = cols

    api_names = _collect_api_names(details)
    if api_names:
        out["api_names"] = api_names

    docs = details.get("documents")
    if isinstance(docs, list) and docs:
        out["document_count"] = len(docs)
        modes = sorted(
            {str(d.get("preview_mode")) for d in docs if isinstance(d, dict) and d.get("preview_mode")}
        )
        if modes:
            out["preview_modes"] = modes

    if dispatch_agent_ids:
        out["dispatch_agent_ids"] = [str(x) for x in dispatch_agent_ids[:12]]
        agents_detail = _dispatch_agents_for_talkback_meta(dispatch_agent_ids)
        if agents_detail:
            out["dispatch_agents"] = agents_detail

    sel = (clarification_selection or "").strip()
    if sel:
        out["clarification_selection"] = sel[:500]

    return out


def _rewrite_context_should_omit_assistant_answer(turn: Dict[str, Any]) -> bool:
    """
    Clarification / disambiguation assistant replies list competing agent names; they must not
    be fed into the contextual question rewriter or the model tends to append those names to
    the user's next question. Preserve the prior user question only (caller adds Q: ...).
    """
    meta = turn.get("answer_meta")
    if not isinstance(meta, dict):
        return False
    if meta.get("needs_clarification"):
        return True
    # Back-compat for rows saved before needs_clarification was persisted on answer_meta
    if str(meta.get("intent", "")).strip().lower() == "ambiguous" and meta.get("success") is False:
        if meta.get("row_count") is None and not meta.get("sql_preview"):
            return True
    return False


def _enriched_answer_for_context(answer: str, answer_meta: Any) -> str:
    """Append a compact machine-readable line for LLM context (relevance / rewrite)."""
    base = str(answer or "").strip()
    if not isinstance(answer_meta, dict) or not answer_meta:
        return base
    bits: List[str] = []
    tt = answer_meta.get("turn_type")
    if isinstance(tt, str) and tt.strip():
        bits.append(f"type={tt.strip()}")
    if answer_meta.get("consent_type"):
        bits.append(f"consent={str(answer_meta.get('consent_type'))[:40]}")
    if answer_meta.get("prior_df_followup") and answer_meta.get("turn_type") == "clarification":
        bits.append(f"prior_df={str(answer_meta.get('prior_df_followup'))[:32]}")
    if answer_meta.get("prior_df_bypass"):
        bits.append("prior_df_llm")
    if "success" in answer_meta:
        bits.append("ok" if answer_meta["success"] else "failed")
    if answer_meta.get("intent"):
        bits.append(f"intent={answer_meta['intent']}")
    if answer_meta.get("datasource"):
        bits.append(f"src={answer_meta['datasource']}")
    if answer_meta.get("row_count") is not None:
        bits.append(f"rows={answer_meta['row_count']}")
    if answer_meta.get("columns"):
        c = answer_meta["columns"]
        if isinstance(c, list) and c:
            bits.append("cols=" + ",".join(c[:12]))
    if answer_meta.get("sql_preview"):
        sp = str(answer_meta["sql_preview"])
        bits.append("sql=" + (sp[:200] + "…" if len(sp) > 200 else sp))
    if answer_meta.get("api_names"):
        an = answer_meta["api_names"]
        if isinstance(an, list) and an:
            bits.append("api=" + ",".join(an[:4]))
    if answer_meta.get("dispatch_agent_ids"):
        da = answer_meta["dispatch_agent_ids"]
        if isinstance(da, list) and da:
            bits.append("agents=" + ",".join(da[:6]))
    if answer_meta.get("talkback_options"):
        tb = answer_meta["talkback_options"]
        if isinstance(tb, list) and tb:
            bits.append(f"tb_opts={len(tb)}")
    if not bits:
        return base
    tag = "[" + " | ".join(bits) + "]"
    return f"{base}\n{tag}" if base else tag


def _turn_for_llm_context(turn: Dict[str, Any]) -> Dict[str, Any]:
    """Copy of a history turn with `answer` augmented by `answer_meta` for classification."""
    t = dict(turn)
    t["answer"] = _enriched_answer_for_context(str(turn.get("answer", "")), turn.get("answer_meta"))
    return t
