agents/supervisor.py

from __future__ import annotations

import json
import re
from typing import Any, Callable

import pandas as pd

from agents.data_extraction import build_data_extraction_node
from agents.data_knowledge import build_data_knowledge_node
from agents.query_analyzer import build_query_analyzer_node
from tools.leading_country_tools import calculate_leading_country, wants_leading_country
from tools.state import GraphState, create_initial_state


LOG_PREFIX = "[Supervisor]"


def _log(msg: str) -> None:
    print(f"{LOG_PREFIX} {msg}")


def _wrap_node(name: str, node_fn: Callable[[GraphState], dict]) -> Callable[[GraphState], dict]:
    """Wrap a graph node to log entry/exit and a one-line summary to the terminal."""

    def wrapped(state: GraphState) -> dict:
        _log(f"→ Entering node: {name}")
        result = node_fn(state)
        summary_parts = []
        if result.get("error_message"):
            err = (result["error_message"] or "")[:60]
            summary_parts.append(f"error={err!r}...")
        if "extracted_data" in result:
            n = len(result.get("extracted_data") or [])
            summary_parts.append(f"extracted_data={n} rows")
        if "parameter_data" in result:
            n = len(result.get("parameter_data") or [])
            summary_parts.append(f"parameter_data={n} rows")
        if "leading_country_result" in result:
            n = len(result.get("leading_country_result") or [])
            summary_parts.append(f"leading_country_result={n} rows")
        if "query_intent" in result:
            summary_parts.append(f"intent={result['query_intent']!r} is_valid={result.get('is_valid', '?')}")
        if "filters" in result and result["filters"]:
            summary_parts.append(f"filters={list(result['filters'].keys())}")
        _log(f"← Exiting node: {name}" + (f" ({', '.join(summary_parts)})" if summary_parts else ""))
        return result

    return wrapped


def _invoke_llm_json(client: Any, prompt: str) -> dict:
    response = client.chat.completions.create(
        model=client._default_deployment,  # type: ignore[attr-defined]
        messages=[
            {"role": "system", "content": "Return strict JSON only."},
            {"role": "user", "content": prompt},
        ],
        temperature=0,
    )
    content = response.choices[0].message.content or "{}"
    try:
        return json.loads(content)
    except Exception:
        match = re.search(r"\{.*\}", content, flags=re.DOTALL)
        if not match:
            return {}
        try:
            return json.loads(match.group(0))
        except Exception:
            return {}


def _invoke_llm_text(client: Any, prompt: str):
    response = client.chat.completions.create(
        model=client._default_deployment,  # type: ignore[attr-defined]
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
    )
    return response.choices[0].message.content or ""


def _format_rows(rows: list[dict], max_rows: int = 10) -> str:
    if not rows:
        return ""
    df = pd.DataFrame(rows[:max_rows])
    return df.to_csv(index=False)


def _deterministic_final_response(state: GraphState) -> str:
    if state.get("error_message"):
        _log("generate_response: returning error_message as final_response")
        return state["error_message"]

    sections: list[str] = []
    extracted_data = state.get("extracted_data", [])
    leading_country_result = state.get("leading_country_result", [])
    parameter_data = state.get("parameter_data", [])

    if extracted_data:
        sections.append(f"Found {len(extracted_data)} matching data rows.")
        sections.append("Sample data:")
        sections.append(f"```csv\n{_format_rows(extracted_data)}\n```")

    if parameter_data:
        sections.append(f"Found {len(parameter_data)} relevant parameter rows.")
        sections.append("Reference data:")
        sections.append(f"```csv\n{_format_rows(parameter_data)}\n```")

    if leading_country_result:
        sections.append(f"Calculated {len(leading_country_result)} leading-country result rows.")
        sections.append("Leading-country calculations:")
        sections.append(f"```csv\n{_format_rows(leading_country_result)}\n```")

    if not sections:
        _log("generate_response: no data → asking user to clarify")
        return "Could you please clarify what you're looking for?"

    sections.append("If you want, I can narrow this further by region, period, product, or market.")
    return "\n\n".join(sections)


def _format_history(history: list[dict], max_turns: int = 6) -> str:
    if not history:
        return ""
    clipped = history[-max_turns:]
    lines: list[str] = []
    for msg in clipped:
        role = str(msg.get("role", "user")).strip().lower()
        content = str(msg.get("content", "")).strip()
        if not content:
            continue
        lines.append(f"{role}: {content}")
    return "\n".join(lines)


def build_generate_response_node(
    llm_text_call: Callable[[str], object] | None = None,
):
    def generate_response_node(state: GraphState) -> dict:
        fallback_response = _deterministic_final_response(state)

        if llm_text_call is None:
            _log("generate_response: llm unavailable, using deterministic response")
            return {"final_response": fallback_response}

        extracted_data = state.get("extracted_data", [])
        leading_country_result = state.get("leading_country_result", [])
        parameter_data = state.get("parameter_data", [])
        filters = state.get("filters", {}) or {}
        conversation_history = state.get("conversation_history", []) or []
        user_query = state.get("user_query", "")

        prompt = f"""
You are the supervisor response writer for a pharma market metrics assistant.
Write the final answer to the user based only on the context below.

Rules:
- Use only facts from the provided context. Do not invent numbers, metrics, or mappings.
- Keep the answer concise and business-friendly.
- Mention applied filters (Region, Period, Calculation_Type, Cluster) briefly.
- When rank columns are present (Current_Volume_Rank, Current_Value_Rank), use them directly
  to determine ordering — the lowest rank number is the best-performing product.
- If the data contains the requested information, answer directly and confidently. Do NOT ask
  clarifying questions when the data is sufficient to answer.
- Only ask a clarifying question if critical information is genuinely missing (e.g. no data returned).
- If Leading-country calculations are present, use them as the source of truth.
- Do not say absolute contribution is unavailable when Current_Value/Current_Volume and Growth fields
  are available; the leading-country calculation reconstructs deltas from those fields.

Context:
User query: {user_query}
Detected intent: {state.get("query_intent", "out_of_scope")}
Error message: {state.get("error_message", "")}
Applied filters: {json.dumps(filters, ensure_ascii=True)}
Extracted data row count: {len(extracted_data)}
Extracted data:
{_format_rows(extracted_data, max_rows=40)}
Leading-country calculation row count: {len(leading_country_result)}
Leading-country calculations:
{_format_rows(leading_country_result, max_rows=40)}
Parameter row count: {len(parameter_data)}
Parameter data sample:
{_format_rows(parameter_data, max_rows=100)}
Recent conversation history:
{_format_history(conversation_history)}

Return plain text only. No markdown code fences.
"""
        try:
            generated = llm_text_call(prompt)
            text = getattr(generated, "content", str(generated)).strip()
            if text:
                _log(
                    "generate_response: llm synthesis success "
                    f"(extracted={len(extracted_data)}, parameter={len(parameter_data)} rows)"
                )
                return {"final_response": text}
        except Exception as exc:
            _log(f"generate_response: llm synthesis failed, using fallback ({exc})")

        _log(
            "generate_response: using deterministic fallback "
            f"(extracted={len(extracted_data)}, parameter={len(parameter_data)} rows)"
        )
        return {"final_response": fallback_response}

    return generate_response_node


def build_leading_country_node(
    db_query_fn: Any | None = None,
    metrics_df: pd.DataFrame | None = None,
):
    def leading_country_node(state: GraphState) -> dict:
        results = calculate_leading_country(
            user_query=state.get("user_query", ""),
            filters=state.get("filters", {}) or {},
            extracted_rows=state.get("extracted_data", []) or [],
            db_query_fn=db_query_fn,
            metrics_df=metrics_df,
        )
        return {"leading_country_result": results}

    return leading_country_node


def route_after_analysis(state: GraphState) -> str:
    is_valid = state.get("is_valid", False)
    intent = state.get("query_intent", "out_of_scope")
    if not is_valid:
        _log(f"route_after_analysis: is_valid=False → generate_response")
        return "generate_response"
    if intent == "data_retrieval":
        _log(f"route_after_analysis: intent={intent!r} → data_extraction")
        return "data_extraction"
    if intent == "parameter_info":
        _log(f"route_after_analysis: intent={intent!r} → data_knowledge")
        return "data_knowledge"
    if intent == "both":
        _log(f"route_after_analysis: intent={intent!r} → data_extraction (then data_knowledge)")
        return "data_extraction"
    _log(f"route_after_analysis: intent={intent!r} → generate_response")
    return "generate_response"


def route_after_data_extraction(state: GraphState) -> str:
    if wants_leading_country(state.get("user_query", "")):
        _log("route_after_data_extraction: leading-country query → leading_country")
        return "leading_country"
    if state.get("query_intent") == "both":
        _log("route_after_data_extraction: intent=both → data_knowledge")
        return "data_knowledge"
    _log("route_after_data_extraction: → generate_response")
    return "generate_response"


def route_after_leading_country(state: GraphState) -> str:
    if state.get("query_intent") == "both":
        _log("route_after_leading_country: intent=both → data_knowledge")
        return "data_knowledge"
    _log("route_after_leading_country: → generate_response")
    return "generate_response"


def build_chatbot_graph(
    metrics_df: pd.DataFrame | None = None,
    azure_openai_client: Any | None = None,
    db_query_fn: Any | None = None,
    column_values: dict | None = None,
):
    """
    Build and compile the LangGraph chatbot.

    ``azure_openai_client`` is named for backward compatibility with the original
    codebase but actually accepts ANY OpenAI-compatible chat client that exposes
    ``client.chat.completions.create(...)`` — including
    ``huggingface_hub.InferenceClient``.
    """
    try:
        from langgraph.graph import END, StateGraph
    except Exception as exc:
        raise ImportError(
            "LangGraph is required. Install with `pip install langgraph`."
        ) from exc

    llm_json = None
    llm_text = None
    if azure_openai_client is not None:
        llm_json = lambda prompt: _invoke_llm_json(azure_openai_client, prompt)
        llm_text = lambda prompt: _invoke_llm_text(azure_openai_client, prompt)

    workflow = StateGraph(GraphState)

    workflow.add_node(
        "query_analyzer",
        _wrap_node("query_analyzer", build_query_analyzer_node(llm_json_call=llm_json)),
    )
    workflow.add_node(
        "data_extraction",
        _wrap_node(
            "data_extraction",
            build_data_extraction_node(
                metrics_df=metrics_df,
                llm_invoke=llm_text,
                db_query_fn=db_query_fn,
                column_values=column_values,
            ),
        ),
    )
    workflow.add_node(
        "data_knowledge",
        _wrap_node("data_knowledge", build_data_knowledge_node()),
    )
    workflow.add_node(
        "leading_country",
        _wrap_node(
            "leading_country",
            build_leading_country_node(db_query_fn=db_query_fn, metrics_df=metrics_df),
        ),
    )
    workflow.add_node(
        "generate_response",
        _wrap_node("generate_response", build_generate_response_node(llm_text_call=llm_text)),
    )

    workflow.set_entry_point("query_analyzer")
    workflow.add_conditional_edges(
        "query_analyzer",
        route_after_analysis,
        {
            "data_extraction": "data_extraction",
            "data_knowledge": "data_knowledge",
            "generate_response": "generate_response",
        },
    )
    workflow.add_conditional_edges(
        "data_extraction",
        route_after_data_extraction,
        {
            "leading_country": "leading_country",
            "data_knowledge": "data_knowledge",
            "generate_response": "generate_response",
        },
    )
    workflow.add_conditional_edges(
        "leading_country",
        route_after_leading_country,
        {
            "data_knowledge": "data_knowledge",
            "generate_response": "generate_response",
        },
    )
    workflow.add_edge("data_knowledge", "generate_response")
    workflow.add_edge("generate_response", END)

    return workflow.compile()


def run_user_query(
    app: Any,
    user_query: str,
    conversation_history: list[dict] | None = None,
) -> GraphState:
    _log(f"User query: {user_query[:80]}{'...' if len(user_query) > 80 else ''}")
    state = create_initial_state(user_query=user_query, conversation_history=conversation_history)
    result = app.invoke(state)
    _log(f"Done. final_response length={len(result.get('final_response', '') or '')} chars")
    return result