src/tools/analytics/merge.py

"""analyze_merge — combine TWO upstream tables on shared keys (KM-608).

The only analytics "family" tool with a SECOND data input. In ONE call it joins
two already-materialized tables (`data` = LEFT, `data_right` = RIGHT) on one or
more shared key columns and returns the combined rows. This is what unlocks the
"which X has BOTH the worst A and the biggest B" question shape: A and B come
from two separate `retrieve_data` pulls (e.g. PA-by-section and backlog-by-section)
that must be aligned per X before either can be judged against the other. Without
a two-input combine the run dies with ColumnNotFoundError because no single tool
ever sees both metrics.

Pattern A, extended: it takes TWO `"${t<id>}"` placeholders. The invoker
materializes BOTH into DataFrames before calling this function (no self-fetch);
the `on` key(s) reference the column aliases the upstream queries produced.

STATUS: compute layer only — takes two already-materialized DataFrames. The
wrapper layer (the ToolOutput envelope, dual-arg materialization, ToolSpec
registration) lives in src/tools/invoker.py + registry.py. Keeping compute
separate from data-fetching keeps this easy to unit-test and stable when wrapped.
"""

from __future__ import annotations

import pandas as pd

from src.tools.analytics.descriptive import ColumnNotFoundError

# Join types the tool understands. Whitelisted so an unknown `how` fails loudly
# instead of silently doing the wrong thing. "cross" is deliberately excluded —
# it ignores `on` and is never the right tool for the "align two metrics" shape.
SUPPORTED_HOWS = ("inner", "left", "right", "outer")


class UnsupportedJoinError(ValueError):
    """Requested join type is not in SUPPORTED_HOWS (maps to error_code UNSUPPORTED_JOIN)."""


def _clean(value: object) -> object:
    """Coerce a scalar to a JSON-clean Python value.

    An outer/left/right join introduces `NaN` for non-matching rows, and numpy /
    pandas scalars (numpy.int64, pandas.Timestamp) are not JSON-serializable —
    normalise all three so the returned rows are clean.
    """
    if isinstance(value, pd.Timestamp):
        return value.isoformat()
    if value is None:
        return None
    try:
        if pd.isna(value):
            return None
    except (TypeError, ValueError):
        pass  # non-scalar / unhashable — leave as-is
    if hasattr(value, "item"):
        return value.item()
    return value


# Prompt-style description read by the Planner to decide WHEN to pick this tool.
DESCRIPTION = """\
Summary: Combine TWO upstream tables into one by joining on shared key column(s) \
(a pandas merge). `data` is the LEFT table, `data_right` is the RIGHT table; `on` \
is the shared column alias(es) present in BOTH. Returns the combined rows, one \
per matched key (join type controlled by `how`, default inner).

USE WHEN a question needs TWO different metrics per the SAME entity and those \
metrics come from two separate pulls — the tell-tale shape is "which X has BOTH \
A and B" (e.g. "which section has the worst PA AND the biggest backlog", "top \
customers by revenue that also have the most complaints"). Plan it as two \
retrieve_data tasks (one per metric, each keyed by X), then analyze_merge on X.

SETTING KEYS: `on` must be column alias(es) that exist in BOTH tables (the entity \
you align on, e.g. section_id). Use `suffixes` (default ["_left","_right"]) to \
disambiguate non-key columns that share a name across the two tables. `how`: \
inner (only matched keys), left/right (keep one side), outer (keep all).

DON'T USE WHEN:
  - both metrics can be pulled in ONE retrieve_data query -> just retrieve_data
  - it groups/aggregates a single table -> analyze_aggregate

Example questions:
  - "which section has the worst PA and the biggest maintenance backlog"
  - "regions in the top 10 for sales that are also bottom 10 for margin"
  - "products low on stock that also have high demand"
"""


def analyze_merge(
    df: pd.DataFrame,
    data_right: pd.DataFrame,
    on: list[str] | str,
    how: str = "inner",
    suffixes: tuple[str, str] | list[str] = ("_left", "_right"),
) -> list[dict[str, object]]:
    """Join two already-materialized tables on shared key column(s).

    Args:
        df: LEFT table (in the real system the invoker materializes this from the
            `data` placeholder).
        data_right: RIGHT table (materialized from the `data_right` placeholder).
        on: shared key column alias(es) present in BOTH tables. A bare string is
            treated as a single key.
        how: join type — one of SUPPORTED_HOWS (default "inner").
        suffixes: 2-element (left, right) suffixes applied to non-key columns that
            collide by name across the two tables.

    Returns:
        list[dict]: one row per merged record, values JSON-clean (NaN -> None).

    Raises:
        ColumnNotFoundError: if `on` is empty or a key is absent from either side.
        UnsupportedJoinError: if `how` is not supported.
        ValueError: if `suffixes` is not a 2-element sequence.
    """
    keys = [on] if isinstance(on, str) else list(on)
    if not keys:
        raise ColumnNotFoundError("merge 'on' must name at least one shared key column")
    if how not in SUPPORTED_HOWS:
        raise UnsupportedJoinError(
            f"unsupported join '{how}'; supported: {list(SUPPORTED_HOWS)}"
        )

    missing_left = [c for c in keys if c not in df.columns]
    missing_right = [c for c in keys if c not in data_right.columns]
    if missing_left or missing_right:
        raise ColumnNotFoundError(
            f"join key(s) not found — left missing {missing_left}, "
            f"right missing {missing_right}"
        )

    suf = tuple(suffixes)
    if len(suf) != 2:
        raise ValueError(f"suffixes must be a 2-element (left, right) sequence, got {suffixes!r}")

    merged = df.merge(data_right, on=keys, how=how, suffixes=suf)
    return [{k: _clean(v) for k, v in rec.items()} for rec in merged.to_dict("records")]