from __future__ import annotations

import hashlib
import json
import logging
import re
import secrets
import time
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Union

from app.services.talkback.types import TalkbackState
from app.services.talkback.utils import _safe_str

logger = logging.getLogger(__name__)

# Pending clarification when context relevance would clear the session; user must confirm first.
OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE = "out_of_context_clear_confirm"

# Must match ``clarification_options[].label`` returned to the client (used for label / action matching).
OUT_OF_CONTEXT_OPTION_LABEL_START_FRESH = "Clear session and continue (fresh routing)"
OUT_OF_CONTEXT_OPTION_LABEL_KEEP_CONTEXT = "Keep session and current agent scope"


def _last_substantive_dispatch_agent_ids(history_turns: Optional[List[Dict[str, Any]]]) -> List[str]:
    """
    ``dispatch_agent_ids`` from the latest history row that reflects an executed answer,
    skipping out-of-context confirmation assistant rows and rows without dispatch metadata.
    """
    for t in reversed(history_turns or []):
        if not isinstance(t, dict):
            continue
        meta = t.get("answer_meta")
        if not isinstance(meta, dict):
            continue
        if meta.get("consent_type") == OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE:
            continue
        da = meta.get("dispatch_agent_ids")
        if not isinstance(da, list) or not da:
            continue
        out = [str(x).strip() for x in da if str(x).strip()]
        if out:
            return out
    return []


def _intent_routing_agent_targets(intent_routing: Optional[Dict[str, Any]]) -> List[str]:
    """Union of routing targets the user may land on (suggested and clarification candidates)."""
    if not isinstance(intent_routing, dict) or intent_routing.get("error"):
        return []
    out: List[str] = []
    seen: set[str] = set()
    for key in ("suggested_agent_ids", "clarification_agent_ids"):
        raw = intent_routing.get(key)
        if not isinstance(raw, list):
            continue
        for x in raw:
            s = str(x).strip()
            if not s or s in seen:
                continue
            seen.add(s)
            out.append(s)
    return out


def _out_of_context_clear_warrants_agent_handoff(
    history_turns: Optional[List[Dict[str, Any]]],
    intent_routing: Optional[Dict[str, Any]],
) -> bool:
    """
    The out-of-context thread confirmation is only needed when the new route targets
    agent(s) disjoint from whoever produced the last substantive answer.
    """
    prior = {str(x).strip() for x in _last_substantive_dispatch_agent_ids(history_turns) if str(x).strip()}
    if not prior:
        return False
    nxt = {str(x).strip() for x in _intent_routing_agent_targets(intent_routing) if str(x).strip()}
    if not nxt:
        return False
    return prior.isdisjoint(nxt)


def _norm_out_of_context_choice_text(text: str) -> str:
    s = str(text or "").strip().lower()
    s = re.sub(r"[\s\u00a0]+", " ", s).strip(" .!?,;")
    for prefix in ("selected:", "selected ", "choice:", "choice ", "option:", "option "):
        if s.startswith(prefix):
            s = s[len(prefix) :].strip()
    return s


def _reply_means_out_of_context_proceed(text: str) -> bool:
    """User chose to clear session and continue with the stored question (fresh thread / new routing)."""
    s = str(text or "").strip().lower()
    if not s:
        return False
    n = _norm_out_of_context_choice_text(text)
    n_label = _norm_out_of_context_choice_text(OUT_OF_CONTEXT_OPTION_LABEL_START_FRESH)
    # Exact / near-exact label from UI, or label embedded in a short phrase.
    if n == n_label or (n_label in n and len(n) <= len(n_label) + 24):
        return True
    # Structured UI: ``clarification_options[].id`` / ``.action``.
    if s in ("start_fresh", "clear_session", "clear_session_proceed", "fresh_routing", "new_thread"):
        return True
    if s == "proceed":
        return True
    if re.match(
        r"^\s*(1|#1|option\s*1|choice\s*1)\s*[.!?]*\s*$",
        s,
    ):
        return True
    compact = re.sub(r"[\s\u00a0]+", " ", s).strip(" .!?,;")
    if "start fresh" in compact or "new thread" in compact or "new agent" in compact:
        return True
    if "clear" in compact and "context" in compact and ("yes" in compact or "ok" in compact):
        return True
    return False


def _reply_means_use_current_context_agents(text: str) -> bool:
    """User chose to keep the talkback session and prior agent scope (no session clear)."""
    s = str(text or "").strip().lower()
    if not s:
        return False
    n = _norm_out_of_context_choice_text(text)
    n_label = _norm_out_of_context_choice_text(OUT_OF_CONTEXT_OPTION_LABEL_KEEP_CONTEXT)
    if n == n_label or (n_label in n and len(n) <= len(n_label) + 24):
        return True
    if s in ("keep_context", "use_current_context", "keep_session", "stay_in_context"):
        return True
    if re.match(
        r"^\s*(2|#2|option\s*2|choice\s*2)\s*[.!?]*\s*$",
        s,
    ):
        return True
    compact = re.sub(r"[\s\u00a0]+", " ", s).strip(" .!?,;")
    if "current context" in compact or "keep context" in compact:
        return True
    if "same agent" in compact or "previous agent" in compact or "prior agent" in compact:
        return True
    if "use current" in compact and ("agent" in compact or "context" in compact):
        return True
    return False


def _history_last_turn_expected_agent_disambiguation_reply(state: TalkbackState) -> bool:
    """True when a recent assistant turn was an agent/source pick clarification (not broad-scope / prior-df)."""
    from app.services.talkback_llm import BROAD_SCOPE_PROCESSING_CONSENT_TYPE

    ht = state.history_turns or []
    for t in reversed(ht):
        if not isinstance(t, dict):
            continue
        q = str(t.get("question", "")).strip()
        if q.lower().startswith("clarification response:"):
            continue
        meta = t.get("answer_meta")
        if not isinstance(meta, dict):
            return False
        ct = str(meta.get("consent_type") or "").strip()
        if ct in (
            BROAD_SCOPE_PROCESSING_CONSENT_TYPE,
            "llm_prior_df",
            "question_clarification",
            OUT_OF_CONTEXT_CLEAR_CONFIRM_CONSENT_TYPE,
        ):
            return False
        if bool(meta.get("needs_clarification")):
            return True
        cids = meta.get("clarification_agent_ids")
        if isinstance(cids, list) and any(str(x).strip() for x in cids):
            return True
        return False
    return False


def _merge_original_question_with_clarification_selection(state: TalkbackState) -> None:
    """
    After the user picks agent(s), restore the ORIGINAL analytical question for execution and history.

    The short reply (e.g. agent names) is kept in ``disambiguation_user_choice`` and in routing metadata,
    but we do **not** append it to ``question`` — ``multi_agent_ask_question`` and talkback session rows
    should store the unmodified analytical question only.
    """
    ht = state.history_turns or []
    if len(ht) < 2:
        return
    last = ht[-1]
    if not isinstance(last, dict):
        return
    last_q = str(last.get("question", "")).strip()
    if not last_q.lower().startswith("clarification response:"):
        return
    choice = (state.question_text or "").strip()
    if not choice:
        return
    orig: Optional[str] = None
    for t in reversed(ht[:-1]):
        if not isinstance(t, dict):
            continue
        tq = str(t.get("question", "")).strip().lower()
        if tq.startswith("clarification response:"):
            continue
        meta = t.get("answer_meta")
        if not isinstance(meta, dict):
            continue
        cids = meta.get("clarification_agent_ids")
        has_disamb = bool(meta.get("needs_clarification")) or (
            isinstance(cids, list) and len(cids) > 0
        )
        if not has_disamb:
            continue
        o = str(t.get("question", "")).strip()
        if o:
            orig = o
            break
    if not orig:
        return
    state.disambiguation_user_choice = choice
    state.talkback_history_question_text = orig
    state.question_text = orig
    try:
        setattr(state.request, "question", orig)
        if getattr(state.request, "my_question", None) is not None:
            setattr(state.request, "my_question", orig)
    except Exception:
        pass


def _last_turn_is_clarification_response_matching_question(state: TalkbackState) -> bool:
    ht = state.history_turns or []
    if not ht:
        return False
    last = ht[-1]
    if not isinstance(last, dict):
        return False
    lq = str(last.get("question", "")).strip()
    if not lq.lower().startswith("clarification response:"):
        return False
    suffix = lq.split(":", 1)[1].strip() if ":" in lq else ""
    return bool(suffix) and suffix == (state.question_text or "").strip()


def _resolve_talkback_effective_session_key(request: Any, user_id: str) -> str:
    """
    Non-empty string for ``derive_talkback_session_key`` third argument and API ``session_key`` echo.

    Reuses stripped ``request.session_key`` when provided; otherwise mints
    ``{uid}::{UTC_ISO8601}::{hex4}`` for a new isolated thread.
    """
    v = getattr(request, "session_key", None)
    if v is not None:
        s = str(v).strip()
        if s:
            return s
    uid = str(user_id or "").strip() or "anonymous"
    ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
    return f"{uid}::{ts}::{secrets.token_hex(4)}"


def talkback_client_session_key_is_empty(request: Any) -> bool:
    """True when the client omitted talkback ``session_key`` or sent only whitespace (new thread)."""
    v = getattr(request, "session_key", None)
    if v is None:
        return True
    return not str(v).strip()


def decide_uploaded_file_followup_action(
    *,
    question_text: str,
    has_file_command: bool,
    has_valid_session_key: bool,
    has_prev_df: bool,
) -> str:
    """
    Decide how talkback should handle question + uploaded file before reconcile routing.

    Returns one of:
      - "clarify_reconcile_vs_analyze": ask user to choose reconcile vs @analyze
      - "auto_analyze": run file analysis directly (no clarification)
      - "default": keep existing reconcile/pipeline behavior
    """
    has_question = bool(str(question_text or "").strip())
    if not has_question or has_file_command:
        return "default"
    if has_valid_session_key and has_prev_df:
        return "clarify_reconcile_vs_analyze"
    if (not has_valid_session_key) and (not has_prev_df):
        return "auto_analyze"
    return "default"


def _derive_talkback_db_session_key(service: Any, state: TalkbackState) -> str:
    """Hashed ``talkback_session_state.session_key`` for this user + effective client ``session_key``."""
    if not (getattr(state, "talkback_effective_session_key", None) or "").strip():
        state.talkback_effective_session_key = _resolve_talkback_effective_session_key(
            state.request, state.current_user_id
        )
    return service.derive_talkback_session_key(
        state.current_user_id,
        state.resolved_agent_ids,
        client_session_key=state.talkback_effective_session_key,
    )


def _agent_scope_list_from_session(session_state: Optional[Dict[str, Any]]) -> Optional[List[str]]:
    """Parse ``agent_scope`` JSON from a session row so we can preserve dispatch scope across early nodes."""
    if not session_state:
        return None
    raw = session_state.get("agent_scope")
    if not raw:
        return None
    try:
        if isinstance(raw, str):
            p = json.loads(raw)
        else:
            p = raw
        if isinstance(p, list):
            return [str(x).strip() for x in p if str(x).strip()]
    except Exception:
        pass
    return None


def _resolve_clarification_via_persisted_agent_scope(
    persisted_scope: Optional[List[str]],
    clarification_agent_ids: Any,
) -> Optional[List[str]]:
    """
    Use the prior turn's dispatch scope (persisted ``agent_scope``) to skip disambiguation
    when the LLM's ``clarification_agent_ids`` still cover all of the previous agents.

    Truth table (prev = persisted_scope, clar = clarification_agent_ids):

      | prev      | clar         | result                                     |
      |-----------|--------------|--------------------------------------------|
      | {1}       | {1}          | [1]   (subset, reuse)                      |
      | {1}       | {1,3}        | [1]   (subset, reuse)                      |
      | {1,5}     | {1,3,5}      | [1,5] (subset, reuse — multi-agent prev)   |
      | {1,5}     | {1,3}        | None  (5 dropped → not subset → ask user)  |
      | {1}       | {3,5}        | None  (1 dropped → not subset → ask user)  |
      | {} or any | empty/invalid| None                                       |

    The previous behavior collapsed multi-agent prev_scope to None and would also pick a
    single overlapping agent from ``clar`` when prev was NOT a subset, which silently dropped
    agents the user had previously chosen. This unified rule is "all-or-nothing": either the
    prior scope is fully recoverable (subset) so we reuse it, or we surface clarification.

    Order in the returned list follows ``persisted_scope`` so dispatch order is stable across
    requests.
    """
    if not clarification_agent_ids or not isinstance(clarification_agent_ids, list):
        return None
    clar_set = {str(a).strip() for a in clarification_agent_ids if str(a).strip()}
    if not clar_set:
        return None
    scope_list = [str(x).strip() for x in (persisted_scope or []) if str(x).strip()]
    if not scope_list:
        return None
    scope_set = set(scope_list)
    if scope_set.issubset(clar_set):
        return scope_list
    return None


def _filter_agent_ids_by_datasource(agent_ids: List[str], datasource: Optional[str]) -> List[str]:
    """Keep only agent_ids that support the selected datasource."""
    ds = str(datasource or "").strip().upper()
    if ds in {"DB", "DATABASE"}:
        target = "DATABASE"
    elif ds == "API":
        target = "API"
    elif ds in {"FILE", "DOCUMENT"}:
        target = "FILE"
    else:
        return list(agent_ids)

    from app.services.agent_config import get_agent_by_id, ordered_available_datasources

    out: List[str] = []
    for aid in agent_ids:
        row = get_agent_by_id(str(aid))
        avail = ordered_available_datasources(row) if row else []
        avail_norm = [str(x).strip().upper() for x in avail]
        if target in avail_norm:
            out.append(str(aid))
    return out


def _infer_talkback_datasource_from_dispatch_and_intent(
    dispatch_agent_ids: List[str],
    intent_value: str,
) -> Optional[str]:
    """
    Choose talkback ``request.datasource`` from agent capabilities first, then intent (IR)
    when more than one capability remains (or when intent picks among compatible options).

    - ``multi_agent`` / ``web``: leave unset (``None``) so downstream mixed / web paths behave as today.
    """
    from app.services.agent_config import get_agent_by_id, ordered_available_datasources

    def _tokens(aid: str) -> List[str]:
        row = get_agent_by_id(str(aid).strip())
        if not row:
            return []
        return [str(x).strip().upper() for x in ordered_available_datasources(row) if str(x).strip()]

    _to_req = {"DATABASE": "Database", "API": "API", "FILE": "File", "DOCUMENT": "File"}

    aids = [str(a).strip() for a in (dispatch_agent_ids or []) if str(a).strip()]
    if not aids:
        return None

    intent_lower = str(intent_value or "").strip().lower()
    if intent_lower in ("multi_agent", "web"):
        return None

    per_agent_sets = [set(_tokens(aid)) for aid in aids]
    if not any(per_agent_sets):
        return None

    supported = per_agent_sets[0] if len(per_agent_sets) == 1 else set.intersection(*per_agent_sets)
    if not supported:
        return None

    def _map(tok: str) -> str:
        return _to_req.get(str(tok).strip().upper(), "Database")

    if len(supported) == 1:
        return _map(next(iter(supported)))

    intent_pref: Optional[str] = None
    if intent_lower == "api":
        intent_pref = "API"
    elif intent_lower == "sql_query":
        intent_pref = "DATABASE"
    elif intent_lower in ("document",):
        intent_pref = "FILE"

    if intent_pref and intent_pref in supported:
        return _map(intent_pref)

    row0 = get_agent_by_id(str(aids[0]))
    for tok in ordered_available_datasources(row0) if row0 else []:
        t = str(tok).strip().upper()
        if t in supported:
            return _map(t)
    return _map(next(iter(supported)))


def _resolve_talkback_file_agent_id(
    state: TalkbackState, candidate_ids: List[str]
) -> Union[str, List[str], None]:
    from app.services.agent_config import resolve_all_file_capable_agent_ids_for_user

    ids = resolve_all_file_capable_agent_ids_for_user(
        str(state.current_user_id or ""),
        candidate_ids,
    )
    if not ids:
        return None
    return ids if len(ids) > 1 else ids[0]


def _state_snapshot(state: TalkbackState) -> Dict[str, Any]:
    prev_q = getattr(state.request, "previous_questions", None)
    prev_count = len(prev_q) if isinstance(prev_q, list) else 0
    return {
        "user_id": _safe_str(state.current_user_id, 48),
        "question_text": _safe_str(state.question_text),
        "datasource": _safe_str(getattr(state.request, "datasource", None), 32),
        "resolved_agent_ids": list(state.resolved_agent_ids),
        "dispatch_agent_ids": list(state.dispatch_agent_ids),
        "history_turns": len(state.history_turns or []),
        "previous_questions": prev_count,
        "context_cleared": state.context_cleared,
        "session_key_set": bool(state.session_key),
        "routing_question_text": _safe_str(state.routing_question_text),
        "intent": _safe_str(state.intent_routing.get("intent"), 32) if isinstance(state.intent_routing, dict) else None,
        "needs_clarification": bool(state.intent_routing.get("needs_clarification"))
        if isinstance(state.intent_routing, dict)
        else None,
        "prior_df_followup": _safe_str(state.intent_routing.get("prior_df_followup"), 24)
        if isinstance(state.intent_routing, dict)
        else None,
        "search_web": bool(getattr(state.request, "search_web", False)),
        "response_set": state.response is not None,
    }


def _print_node(tag: str, node_name: str, trace_id: str, state: TalkbackState, **extra: Any) -> None:
    base = {
        "trace_id": trace_id,
        "node": node_name,
        "tag": tag,
        "state": _state_snapshot(state),
    }
    if extra:
        base["params"] = {k: v for k, v in extra.items()}
    print(f"[TalkbackPipeline] {json.dumps(base, ensure_ascii=False)}")


def _build_intent_routing_question(state: TalkbackState) -> str:
    """
    Include the last Q&A in the routing prompt when context was kept, so the model can set
    prior_df_followup (e.g. forecast / analysis on the previous result → askme_analysis).

    When there is history, use the **verbatim** user message for "Current user message", not
    `routing_question_text`. The rewrite step (`ContextRewriteForRoutingNode`) expands shorthand
    into a standalone sentence (e.g. naming the Order Management system) which often reads
    like a **new** sql_query to the router and incorrectly sets prior_df_followup to none.
    For the first turn (no history), keep using the rewrite so domain expansion still helps.
    """
    if state.context_cleared or not state.history_turns:
        return (state.routing_question_text or state.question_text or "").strip()
    base = (state.question_text or "").strip()
    last = state.history_turns[-1]
    if not isinstance(last, dict):
        return base
    last_q = str(last.get("question", "")).strip()

    if last_q.lower().startswith("clarification response:"):
        pick = getattr(state, "disambiguation_user_choice", None)
        if isinstance(pick, str) and pick.strip():
            base = pick.strip()
        # The current message is a follow-up to a disambiguation prompt.
        # Walk back to find the original clarification turn that has answer_meta
        # so the router knows *which agents* were competing and what the user chose.
        for t in reversed(state.history_turns[:-1]):
            if not isinstance(t, dict):
                continue
            meta = t.get("answer_meta")
            if not isinstance(meta, dict):
                continue
            cids = meta.get("clarification_agent_ids")
            has_disamb = bool(meta.get("needs_clarification")) or (
                isinstance(cids, list) and len(cids) > 0
            )
            if not has_disamb:
                continue
            clar_msg = str(t.get("answer", "")).strip()[:400]
            orig_q = str(t.get("question", "")).strip()[:500]
            cids_info = ""
            if isinstance(cids, list) and cids:
                cids_info = f"\nCompeting agent_ids: {', '.join(str(x) for x in cids)}"
            return (
                f"The system previously asked the user to disambiguate between agents:\n"
                f"Original question: {orig_q}\n"
                f"System clarification prompt: {clar_msg}{cids_info}\n\n"
                f"The user's response (choosing agent(s)):\n{base}\n\n"
                f"IMPORTANT: If the user says 'both', 'all', 'from both agents', or names "
                f"specific agents, set intent to 'multi_agent' with those agent ids in "
                f"suggested_agent_ids. Do NOT ask for clarification again."
            )
        return base

    last_a = str(last.get("answer", "")).strip()
    last_q = last_q[:500]
    last_a = last_a[:400]
    if not last_q:
        return base
    return (
        "Prior conversation turn (follow-ups may need the previous tabular result for @askme / @web):\n"
        f"Q: {last_q}\n"
        f"A (preview): {last_a}\n\n"
        f"Current user message:\n{base}"
    )


def _derive_request_id(state: TalkbackState) -> str:
    """
    Idempotency key for the current talkback request.
    - If client provided `request_id`, use it.
    - Otherwise derive deterministically from identity + question + current agent scope.
    """
    client_rid = getattr(state.request, "request_id", None)
    if client_rid is not None:
        rid = str(client_rid).strip()
        if rid:
            return rid

    agent_ids = state.dispatch_agent_ids or state.resolved_agent_ids or []
    raw = f"{state.current_user_id}::{'|'.join(sorted(map(str, agent_ids)))}::{state.question_text}::{state.context_cleared}"
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()


def _derive_request_id_for_direct(
    request: Any, current_user_id: str, resolved_agent_ids: List[str], question_text: str
) -> str:
    client_rid = getattr(request, "request_id", None)
    if client_rid is not None:
        rid = str(client_rid).strip()
        if rid:
            return rid
    raw = f"{current_user_id}::{'|'.join(sorted(map(str, resolved_agent_ids)))}::{question_text}::False"
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()
