core/bind.py

from __future__ import annotations
import json
from pathlib import Path

from core.ports import ChatModel
from core.schemas import BindResponse, PickVarsResponse
from core import runner


# ---------------------------------------------------------------------------
# Sheet detection
# ---------------------------------------------------------------------------

def _match_sheet(sheets: list[str], keyword: str) -> str:
    """Return the sheet name that best fuzzy-matches keyword (no external deps)."""
    def score(name: str) -> int:
        n = name.lower()
        if n == keyword:
            return 3
        if keyword in n:
            return 2
        if n.startswith(keyword[0]):
            return 1
        return 0
    return max(sheets, key=score, default=keyword)


def suggest_sheets(sheets: list[str]) -> tuple[str, str]:
    """Return (survey_sheet, choices_sheet) best-guess names from a sheet list."""
    return _match_sheet(sheets, "survey"), _match_sheet(sheets, "choices")


def list_kobo_sheets(xlsx_path: Path) -> list[str]:
    """Call read_kobo.py --list-sheets; return sheet names as a list.

    The script returns {"source_file": ..., "sheet_names": [...], ...}.
    """
    result = runner.run_skill_script("read_kobo.py", [str(xlsx_path), "--list-sheets"])
    if isinstance(result, dict):
        for key in ("sheet_names", "sheets"):
            if key in result and isinstance(result[key], list):
                return list(result[key])
        for v in result.values():
            if isinstance(v, list):
                return list(v)
        return []
    return list(result)


# ---------------------------------------------------------------------------
# Kobo parsing
# ---------------------------------------------------------------------------

def parse_kobo(
    xlsx_path: Path,
    slug: str,
    out_path: Path,
    survey_sheet: str = "survey",
    choices_sheet: str = "choices",
) -> Path:
    """Parse an XLSForm into a kobo cache JSON file.

    Calls read_kobo.py in phase-2 mode (--slug). Writes to out_path.
    Returns out_path on success.
    """
    runner.run_skill_script_raw(
        "read_kobo.py",
        [
            str(xlsx_path),
            "--slug", slug,
            "-o", str(out_path),
            "--survey-sheet", survey_sheet,
            "--choices-sheet", choices_sheet,
        ],
    )
    return out_path


# ---------------------------------------------------------------------------
# Kobo query helpers
# ---------------------------------------------------------------------------

def kobo_summary(cache_path: Path) -> dict:
    """Return the cheap orientation summary (~200B name map) from a kobo cache."""
    return runner.run_skill_script(
        "read_kobo.py",
        ["--cache", str(cache_path), "--summary"],
    )


def kobo_names(cache_path: Path, names: list[str]) -> dict:
    """Fetch full variable details (with choices) for the given variable names.

    Choices are included by default; read_kobo.py only accepts --no-choices to omit
    them, so we pass no choices flag at all.
    """
    return runner.run_skill_script(
        "read_kobo.py",
        ["--cache", str(cache_path), "--names", ",".join(names)],
    )


# ---------------------------------------------------------------------------
# Candidate-variable normalisation
# ---------------------------------------------------------------------------

def normalise_candidate_vars(candidates: list[str], labels_map: dict) -> list[str]:
    """Coerce model-proposed candidates to valid survey variable names.

    The model is asked for variable names but sometimes returns a question label or a
    ``name: label`` pair. We resolve each candidate against the survey's known names
    (the keys of ``labels_map``) and drop anything that cannot be resolved:

    1. exact variable name → keep
    2. ``name: label`` → take the part before the first colon if it is a known name
    3. a bare question label → reverse-map to its variable name
    4. otherwise → drop (better to show no variable than a wrong label)

    Order and de-duplication are preserved.
    """
    valid_names = set(labels_map.keys())
    # Reverse map: label text → variable name (first wins on duplicate labels).
    label_to_name: dict[str, str] = {}
    for name, lbl in labels_map.items():
        label_to_name.setdefault(str(lbl).strip(), name)

    resolved: list[str] = []
    for raw in candidates:
        cand = str(raw).strip()
        name = None
        if cand in valid_names:
            name = cand
        elif ":" in cand and cand.split(":", 1)[0].strip() in valid_names:
            name = cand.split(":", 1)[0].strip()
        elif cand in label_to_name:
            name = label_to_name[cand]
        if name is not None and name not in resolved:
            resolved.append(name)
    return resolved


# ---------------------------------------------------------------------------
# Bind loop — two LLM calls per indicator
# ---------------------------------------------------------------------------

def run_bind_step(
    indicator_id: str,
    indicator_def: dict,
    cache_path: Path,
    summary_map: dict,
    chat_model: ChatModel,
) -> BindResponse:
    """Bind one indicator to the dataset via two LLM calls.

    Call 1: summary_map + definition → PickVarsResponse (1–2 candidate var names).
    Call 2: var details + definition + errors + ki_note → BindResponse (verdict).
    """
    label = indicator_def.get("label", indicator_id)
    definition = indicator_def.get("definition", "")
    errors = indicator_def.get("common_implementation_errors", "—")
    ki_note = indicator_def.get("ki_assessment_note", "—")

    if isinstance(summary_map, dict) and summary_map.get("question_labels"):
        labels_map = summary_map["question_labels"]
    elif isinstance(summary_map, dict) and summary_map.get("all_question_names"):
        labels_map = {n: n for n in summary_map["all_question_names"]}
    else:
        # legacy flat {name: description} map (older callers / fixtures)
        labels_map = summary_map if isinstance(summary_map, dict) else {}

    summary_text = "\n".join(f"- {name}: {lbl}" for name, lbl in labels_map.items())
    pick_messages = [
        {
            "role": "user",
            "content": (
                f"You are mapping humanitarian indicators to a Kobo survey instrument.\n\n"
                f"Indicator: {label}\n"
                f"Definition: {definition}\n\n"
                f"Survey variables (variable_name: question label):\n{summary_text}\n\n"
                f"Pick 1–2 variable_names from the list whose question label best matches "
                f"this indicator. Return your answer as JSON using the variable_name (the part "
                f"before the colon). If no variable matches, return an empty list."
            ),
        }
    ]
    try:
        pick: PickVarsResponse = chat_model.structured(pick_messages, PickVarsResponse)
    except Exception:
        # No raw-error leak (PICK): this prompt feeds every survey label (~20k tokens) and is the
        # call that truncated in the field. A parse/length failure here must not propagate to the
        # app (which would write "Bind error: …" into the spec). Fall back deterministically.
        return BindResponse(
            indicator_id=indicator_id,
            variables=[],
            measurable="NOT_MEASURABLE",
            reasons=(
                "Bind inconclusive: variable-selection response could not be parsed "
                "(likely truncated output); defaulting to NOT_MEASURABLE. "
                "Re-run or review manually."
            ),
            result_ids=[],
        )
    # Code disposes: resolve the model's proposals to valid survey variable names and
    # treat THIS as the binding's variable list. Call 2 only classifies the verdict; we
    # never trust the verdict call's re-emitted variable list (it leaks question labels).
    candidate_vars = normalise_candidate_vars(pick.candidate_variables, labels_map)

    # force-NONE guard: if no proposal resolved to a real survey variable, CODE (not the model)
    # decides the verdict. The model is never given the chance to claim MEASURABLE on a
    # fabricated or empty match — the false-positive pattern. Skips the verdict call entirely.
    if not candidate_vars:
        return BindResponse(
            indicator_id=indicator_id,
            variables=[],
            measurable="NOT_MEASURABLE",
            reasons=(
                "No survey variable maps to this indicator "
                "(no candidate passed the variable-name allowlist)."
            ),
            result_ids=[],
        )

    var_details: dict = kobo_names(cache_path, candidate_vars)
    details_text = (
        json.dumps(var_details, ensure_ascii=False, indent=2)
        if var_details
        else "(no matching variable found)"
    )
    verdict_messages = [
        {
            "role": "user",
            "content": (
                f"You are deciding whether a Kobo survey can measure a humanitarian indicator.\n\n"
                f"Indicator: {label}\n"
                f"Definition: {definition}\n"
                f"Common errors: {errors}\n"
                f"KI assessment note: {ki_note}\n\n"
                f"Candidate variable(s) from the survey:\n{details_text}\n\n"
                f"Decide: MEASURABLE | PROXY | NOT_MEASURABLE.\n"
                f"MEASURABLE: exact construct, correct unit of analysis, answer options map without transformation.\n"
                f"PROXY: same construct but different format / unit / missing criteria.\n"
                f"NOT_MEASURABLE: no matching variable or instrument cannot compute the indicator.\n\n"
                f"In the reasons field, state what the binding proves AND what it cannot prove."
            ),
        }
    ]
    try:
        bind: BindResponse = chat_model.structured(verdict_messages, BindResponse)
    except Exception:
        # No raw-error leak: a truncated / un-parseable verdict (e.g. the model rambling to
        # max_tokens) must never dump exception text into the spec's reasons field. Fall back
        # conservatively to NOT_MEASURABLE, keeping the variables we did resolve. Conservative
        # bias (never a false MEASURABLE) is the right default for a humanitarian tool.
        return BindResponse(
            indicator_id=indicator_id,
            variables=candidate_vars,
            measurable="NOT_MEASURABLE",
            reasons=(
                "Bind inconclusive: the model verdict could not be parsed "
                "(likely truncated output); defaulting to NOT_MEASURABLE. "
                "Re-run or review manually."
            ),
            result_ids=[],
        )
    return BindResponse(
        indicator_id=indicator_id,
        variables=candidate_vars,
        measurable=bind.measurable,
        reasons=bind.reasons,
        result_ids=bind.result_ids,
    )