src/agent/nodes.py

"""
LangGraph nodes.

Flow:
  START → orchestrator ─(has tool_calls)→ tool_node → orchestrator (loop)
                       ↘ (no tool_calls) → masker → response → END

masker_node:
  Reads all ToolMessages in state["messages"].
  - ToolMessages from search_graph  → stored in graph_data (contains copyrighted defs)
  - ToolMessages from wikipedia / tavily_search → stored in external_defs
  Replaces every concept's "definition" field with an external one, or removes it.
  Result goes into masked_data, which response_node uses as context.
"""
from __future__ import annotations
import json
from langchain_core.messages import SystemMessage, ToolMessage, AIMessage, HumanMessage
from langgraph.prebuilt import ToolNode

from src.llm import llm
from src.agent.state import AgentState
from src.agent.prompts import ORCHESTRATOR_SYSTEM, RESPONSE_SYSTEM
from src.tools.graph_tool import ALL_GRAPH_TOOLS
from src.tools.wikipedia_tool import wikipedia
from src.tools.tavily_tool import tavily

ALL_TOOLS = ALL_GRAPH_TOOLS + [wikipedia, tavily]
tool_node = ToolNode(ALL_TOOLS, handle_tool_errors=lambda e: f"[Tool error — {type(e).__name__}] {e}")
_llm_with_tools = llm.bind_tools(ALL_TOOLS)

_MAX_MESSAGES_BEFORE_SUMMARY = 30
_MESSAGES_TO_KEEP = 16


def _plain_ai(msg: AIMessage) -> AIMessage:
    """
    Return a brand-new AIMessage containing only the plain text content.
    Strips additional_kwargs, response_metadata, tool_calls and any other
    SDK-specific attributes that langchain-google-genai may have attached.
    When an AIMessage from a previous turn is sent back as history, those
    attributes can cause Gemini to misclassify the turn as a function-call
    turn and raise INVALID_ARGUMENT.
    """
    content = getattr(msg, "content", "") or ""
    if isinstance(content, list):
        content = " ".join(
            p.get("text", "") if isinstance(p, dict) else str(p)
            for p in content
        ).strip()
    return AIMessage(content=str(content))


def _messages_for_response(messages: list) -> list:
    """
    Build a clean H→AI→H→AI sequence for response_node.
    Keeps HumanMessages + final-response AIMessages (those followed by a
    HumanMessage). Past AIMessages are rebuilt via _plain_ai() to strip any
    SDK metadata that could confuse Gemini's turn-validation.
    """
    result = []
    for i, msg in enumerate(messages):
        if isinstance(msg, HumanMessage):
            result.append(msg)
        elif isinstance(msg, AIMessage) and not getattr(msg, "tool_calls", None):
            nxt = messages[i + 1] if i + 1 < len(messages) else None
            if isinstance(nxt, HumanMessage):
                result.append(_plain_ai(msg))
    return result


def _messages_for_orchestrator(messages: list) -> list:
    """
    Build the message list for orchestrator_node.
    Previous turns: clean H→AI pairs (AIMessages rebuilt via _plain_ai).
    Current turn (from last HumanMessage): full tool-call context unchanged.
    """
    current_start = 0
    for i in range(len(messages) - 1, -1, -1):
        if isinstance(messages[i], HumanMessage):
            current_start = i
            break

    previous: list = []
    if current_start > 0:
        prev = list(messages[:current_start])
        for i, msg in enumerate(prev):
            if isinstance(msg, HumanMessage):
                previous.append(msg)
            elif isinstance(msg, AIMessage) and not getattr(msg, "tool_calls", None):
                nxt = prev[i + 1] if i + 1 < len(prev) else None
                if isinstance(nxt, HumanMessage) or nxt is None:
                    previous.append(_plain_ai(msg))

    return previous + list(messages[current_start:])


async def _maybe_summarize(state: AgentState) -> tuple[list, str | None]:
    """
    If history exceeds _MAX_MESSAGES_BEFORE_SUMMARY, summarize the older portion.
    Returns (messages_to_use, new_summary_or_None).
    """
    messages = list(state["messages"])
    if len(messages) <= _MAX_MESSAGES_BEFORE_SUMMARY:
        return messages, None

    old_msgs = messages[:-_MESSAGES_TO_KEEP]
    recent_msgs = messages[-_MESSAGES_TO_KEEP:]

    lines = [
        f"[{getattr(m, 'type', '?')}]: {str(getattr(m, 'content', ''))[:300]}"
        for m in old_msgs
        if getattr(m, "content", None)
    ]
    history_text = "\n".join(lines)

    prior = state.get("conversation_summary") or ""
    if prior:
        history_text = f"Tóm tắt cũ:\n{prior}\n\nNội dung bổ sung:\n{history_text}"

    result = await llm.ainvoke([
        SystemMessage(content=(
            "Tóm tắt ngắn gọn cuộc hội thoại sau bằng tiếng Việt. "
            "Giữ nguyên thuật ngữ tiếng Anh. "
            "Chỉ giữ những gì người dùng đã hỏi và thông tin quan trọng đã cung cấp."
        )),
        HumanMessage(content=history_text),
    ])
    return recent_msgs, result.content


# ─────────────────────────────────────────────────────────────────────────────
# Orchestrator
# ─────────────────────────────────────────────────────────────────────────────

async def orchestrator_node(state: AgentState) -> dict:
    context_messages, new_summary = await _maybe_summarize(state)

    summary = new_summary or state.get("conversation_summary") or ""
    system_content = ORCHESTRATOR_SYSTEM
    if summary:
        system_content += f"\n\n## TÓM TẮT HỘI THOẠI TRƯỚC:\n{summary}"

    messages = [SystemMessage(content=system_content)] + _messages_for_orchestrator(context_messages)
    response = await _llm_with_tools.ainvoke(messages)

    for tc in getattr(response, "tool_calls", []):
        print(f"[tool] {tc['name']}  args={tc.get('args', {})}", flush=True)

    updates: dict = {"messages": [response]}
    if new_summary:
        updates["conversation_summary"] = new_summary
    return updates


_MAX_TOOL_ITERS = 5


def should_continue(state: AgentState) -> str:
    last = state["messages"][-1]
    if not (isinstance(last, AIMessage) and getattr(last, "tool_calls", None)):
        return "masker"

    # Count tool-call rounds in the current turn (since last HumanMessage)
    iters = 0
    for msg in reversed(state["messages"]):
        if isinstance(msg, HumanMessage):
            break
        if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None):
            iters += 1

    return "tools" if iters < _MAX_TOOL_ITERS else "masker"


# ─────────────────────────────────────────────────────────────────────────────
# Masker
# ─────────────────────────────────────────────────────────────────────────────

def masker_node(state: AgentState) -> dict:
    # Build tool_call_id → {name, args} map from all AIMessages
    call_map: dict[str, dict] = {}
    for msg in state["messages"]:
        if isinstance(msg, AIMessage):
            for tc in getattr(msg, "tool_calls", []):
                call_map[tc["id"]] = {"name": tc["name"], "args": tc.get("args", {})}

    graph_data: dict = {}
    external_defs: dict[str, str] = {}

    for msg in state["messages"]:
        if not isinstance(msg, ToolMessage):
            continue
        tc_info  = call_map.get(msg.tool_call_id, {})
        tool_name = tc_info.get("name", "")
        args      = tc_info.get("args", {})
        content   = msg.content

        _graph_tool_names = {t.name for t in ALL_GRAPH_TOOLS}
        if tool_name in _graph_tool_names:
            try:
                graph_data.update(json.loads(content))
            except (json.JSONDecodeError, TypeError):
                pass

        elif tool_name in ("wikipedia", "tavily_search"):
            # Derive the concept name from the tool's input argument
            concept_key = (
                args.get("query")
                or args.get("tool_input")
                or args.get("__arg1")
                or f"__{tool_name}_{msg.tool_call_id}"
            )
            if isinstance(content, str):
                external_defs[concept_key] = content
            elif isinstance(content, list):
                # Tavily returns list of dicts with "content" key
                texts = [
                    item.get("content", "") or item.get("answer", "")
                    for item in content
                    if isinstance(item, dict)
                ]
                external_defs[concept_key] = " ".join(t for t in texts if t)

    # Apply masking: replace or strip the "definition" field on every concept
    def _mask_concept_list(concepts: list[dict]) -> list[dict]:
        out = []
        for c in concepts:
            masked = {k: v for k, v in c.items() if k != "definition"}
            ext = _best_external_def(c.get("name", ""), external_defs)
            if ext:
                masked["definition"] = ext
            out.append(masked)
        return out

    masked: dict = {}
    for key, val in graph_data.items():
        if isinstance(val, list) and val and isinstance(val[0], dict) and "definition" in val[0]:
            masked[key] = _mask_concept_list(val)
        else:
            masked[key] = val

    return {
        "graph_data":    graph_data,
        "external_defs": external_defs,
        "masked_data":   masked,
    }


def _best_external_def(concept_name: str, external_defs: dict[str, str]) -> str | None:
    """
    Find the external definition most relevant to concept_name.
    Strategy: prefer exact key match, then substring, then first-sentence keyword scan.
    """
    name_lower = concept_name.lower()

    # 1. Exact key match
    for key, text in external_defs.items():
        if key.lower() == name_lower:
            return _first_relevant_sentence(concept_name, text)

    # 2. Key contains concept name or vice-versa
    for key, text in external_defs.items():
        if name_lower in key.lower() or key.lower() in name_lower:
            return _first_relevant_sentence(concept_name, text)

    # 3. Any external text that contains the concept name
    for text in external_defs.values():
        if name_lower in text.lower():
            return _first_relevant_sentence(concept_name, text)

    return None


def _first_relevant_sentence(concept_name: str, text: str) -> str | None:
    name_lower = concept_name.lower()
    for sentence in text.replace("\n", " ").split("."):
        sentence = sentence.strip()
        if sentence and name_lower in sentence.lower():
            return sentence + "."
    # Fallback: first non-empty sentence of the text
    for sentence in text.replace("\n", " ").split("."):
        sentence = sentence.strip()
        if len(sentence) > 20:
            return sentence + "."
    return None


# ─────────────────────────────────────────────────────────────────────────────
# Response
# ─────────────────────────────────────────────────────────────────────────────

async def response_node(state: AgentState) -> dict:
    from datetime import datetime, timezone, timedelta
    tz_vn = timezone(timedelta(hours=7))
    now = datetime.now(tz_vn).strftime("%A, %d/%m/%Y %H:%M UTC+7")

    masked_data = state.get("masked_data", {})

    # Always extract orchestrator's final text from the current turn.
    # When no tools were called: becomes primary context.
    # When tools were called: added as supplementary reasoning alongside masked tool data.
    for msg in reversed(state["messages"]):
        if isinstance(msg, HumanMessage):
            break
        if isinstance(msg, AIMessage) and not getattr(msg, "tool_calls", None):
            content = msg.content
            if isinstance(content, list):
                content = " ".join(
                    p.get("text", "") if isinstance(p, dict) else str(p)
                    for p in content
                ).strip()
            text = str(content or "")
            if text:
                if not masked_data:
                    masked_data = {"thông_tin": text}
                else:
                    masked_data = {**masked_data, "ghi_chú_orchestrator": text}
            break

    context = json.dumps(masked_data, ensure_ascii=False, indent=2)
    system = RESPONSE_SYSTEM.format(context=context, current_datetime=now)

    summary = state.get("conversation_summary") or ""
    if summary:
        system += f"\n\n## TÓM TẮT HỘI THOẠI TRƯỚC:\n{summary}"

    history = _messages_for_response(list(state["messages"]))
    print(f"[response_node] history={[type(m).__name__ + '(tc=' + str(bool(getattr(m,'tool_calls',None))) + ')' for m in history]}", flush=True)
    messages = [SystemMessage(content=system)] + history
    response = await llm.ainvoke(messages)
    return {"messages": [response]}