src/tools/dataset_tools/_base.py

"""Shared base for the dataset_tools package: imports, dataset_mcp, helpers, intent constants."""
from __future__ import annotations
# ruff: noqa: F401, E402  (imports re-exported to tool submodules; intentional late imports)
import sys
from pathlib import Path
from typing import Annotated, Any, Optional

import pandas as pd
from fastmcp import FastMCP

_PROJECT_ROOT = Path(__file__).parent.parent.parent.resolve()
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

from src.datasets.registry import load_manifest, list_available_datasets, get_integration_plan
from src.datasets.manifest_schema import validate_manifest
from src.workflows.microarray import get_collapse_params, recommend_analysis_path
from src.workflows.activity_stats import get_contrast_groups
from src.workflows.metadata_validation import list_valid_groups, validate_contrast

dataset_mcp = FastMCP(name="dataset")


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _load_metadata_df(metadata_path: str) -> pd.DataFrame:
    """
    Load sample metadata into a DataFrame from .h5ad, .csv, or .tsv.

    For .h5ad files, returns adata.obs (sample-level metadata only).
    Raises FileNotFoundError if the path does not exist.
    """
    path = Path(metadata_path)
    if not path.exists():
        raise FileNotFoundError(f"Metadata file not found: {metadata_path}")

    suffix = path.suffix.lower()
    if suffix == ".h5ad":
        import scanpy as sc
        adata = sc.read_h5ad(str(path))
        return adata.obs.copy()
    elif suffix in (".tsv", ".txt"):
        return pd.read_csv(str(path), sep="\t", index_col=0)
    else:
        return pd.read_csv(str(path), index_col=0)


def _build_loading_plan(manifest: Any) -> list[dict]:
    """Build a step-by-step loading plan from a DatasetManifest."""
    source = manifest.expression_source
    meta_source = manifest.metadata_source
    collapse = get_collapse_params(manifest)
    contrast = get_contrast_groups(manifest)
    pid = manifest.dataset_id

    steps = []
    step = 1
    collapse_precomputed = False

    # Single-cell / spatial (Path P) is decided by MODALITY, not the source-type
    # string: the expression_source is a hosted AnnData h5ad whether the manifest
    # declares it as `h5ad` or as `url` pointing at a `.h5ad` (biodata-registry
    # uses `url`). Load via the sc loader (read_h5ad_cached, ADR-0006 Role 1) —
    # NOT the bulk flat-file loader `decoupler_load_url_counts`.
    if manifest.analysis_path == "P":
        steps.append({
            "step": step,
            "tool": "decoupler_load_and_visualize_data",
            "key_args": {
                "adata_path": source.get("url"),
                "out_prefix": f"{pid}_loaded",
            },
            "note": (
                "Single-cell / spatial h5ad (ADR-0006 Role 1). Loads the hosted "
                "AnnData via read_h5ad_cached — parsed once per resident MCP "
                "process; a hosted/private h5ad URL is resolved to a local path "
                "via the authenticated loader (HF_TOKEN). The full integrated "
                "atlas is NOT served on this live path — score its offline-derived "
                "signatures against bulk cohorts with dataset_score_signature."
            ),
        })
        step += 1
        expr_h5ad = f"<output_path from step {step - 1}>"

    elif source.get("type") == "geo_series_matrix":
        steps.append({
            "step": step,
            "tool": "decoupler_load_geo_series_matrix",
            "key_args": {
                "url_or_path": source.get("url"),
                "condition_column": manifest.group_columns[0] if manifest.group_columns else None,
                "out_prefix": f"{pid}_raw",
            },
            "note": "Downloads GEO series matrix, decodes numeric condition codes, saves h5ad.",
        })
        step += 1
        expr_h5ad = f"<output_path from step {step - 1}>"

    elif source.get("type") == "url":
        collapsed_url = source.get("collapsed_url")
        collapse_precomputed = bool(collapsed_url) and collapse["required"]
        if collapse_precomputed:
            load_note = (
                f"Downloads precomputed gene-collapsed expression "
                f"(probes already collapsed via {collapse['method']}); "
                f"no separate decoupler_collapse_probes_to_genes step needed."
            )
        else:
            load_note = (
                f"Downloads flat genes × samples matrix (URL type). "
                f"feature_id_type='{manifest.feature_id_type}'. "
                f"Transposes to AnnData convention (samples × genes)."
            )
        steps.append({
            "step": step,
            "tool": "decoupler_load_url_counts",
            "key_args": {
                "url_or_path": collapsed_url if collapse_precomputed else source.get("url"),
                "feature_id_type": manifest.feature_id_type,
                "strip_ensembl_versions": manifest.feature_id_type == "ensembl_gene_id",
                "out_prefix": f"{pid}_raw",
            },
            "note": load_note,
        })
        step += 1
        expr_h5ad = f"<output_path from step {step - 1}>"

        # Clinical join: only when metadata_source has a URL (not embedded)
        clin_url = meta_source.get("url") if meta_source else None
        if clin_url and not meta_source.get("embedded", False):
            steps.append({
                "step": step,
                "tool": "decoupler_join_clinical_metadata",
                "key_args": {
                    "adata_path": expr_h5ad,
                    "clinical_url_or_path": clin_url,
                    "barcode_column": meta_source.get("join_column"),
                    "truncate_to_patient": meta_source.get("truncate_to_patient", True),
                    "add_survival_columns": bool(manifest.survival_columns.get("event_column")),
                    "out_prefix": f"{pid}_clinical",
                },
                "note": (
                    "Joins clinical/phenotype TSV to the expression h5ad by sample barcode. "
                    "Derives os_event/os_days from vital_status + days_to_death if survival columns are defined."
                ),
            })
            step += 1
            expr_h5ad = f"<output_path from step {step - 1}>"

        # Curation step: filter to curated sample list if defined in manifest
        if getattr(manifest, "curated_sample_list", None):
            n_curated = len(manifest.curated_sample_list)
            steps.append({
                "step": step,
                "tool": "dataset_filter_to_curated_samples",
                "key_args": {
                    "adata_path": expr_h5ad,
                    "dataset_id": pid,
                    "out_prefix": f"{pid}_curated",
                },
                "note": (
                    f"Filters to {n_curated} curated samples from the manifest "
                    f"({getattr(manifest, 'curated_sample_source', 'see manifest')}). "
                    "Required before survival or prognostic analysis to remove "
                    "non-target tissue contamination."
                ),
            })
            step += 1
            expr_h5ad = f"<output_path from step {step - 1}>"

    elif source.get("type") == "gdc":
        steps.append({
            "step": step,
            "tool": "decoupler_load_gdc_star_counts",
            "key_args": {
                "project_id": source.get("project_id", "TCGA-PAAD"),
                "count_column": source.get("count_column", "unstranded"),
                "feature_id_type": manifest.feature_id_type,
                "strip_ensembl_versions": manifest.feature_id_type == "ensembl_gene_id",
                "out_prefix": f"{pid}_raw",
            },
            "note": (
                f"Downloads STAR-Counts from GDC API for {source.get('project_id', 'TCGA-PAAD')}. "
                "Bulk-downloads ~178 per-sample TSVs as a single tar.gz (~500 MB). "
                "Returns integer raw counts — Path A (DESeq2). No authentication required."
            ),
        })
        step += 1
        expr_h5ad = f"<output_path from step {step - 1}>"

        # Clinical join: only when metadata_source has a URL (not embedded)
        clin_url = meta_source.get("url") if meta_source else None
        if clin_url and not meta_source.get("embedded", False):
            steps.append({
                "step": step,
                "tool": "decoupler_join_clinical_metadata",
                "key_args": {
                    "adata_path": expr_h5ad,
                    "clinical_url_or_path": clin_url,
                    "barcode_column": meta_source.get("join_column"),
                    "truncate_to_patient": meta_source.get("truncate_to_patient", True),
                    "add_survival_columns": bool(manifest.survival_columns.get("event_column")),
                    "out_prefix": f"{pid}_clinical",
                },
                "note": (
                    "Joins clinical/phenotype TSV to the expression h5ad by sample barcode. "
                    "Derives os_event/os_days from vital_status + days_to_death if survival columns are defined."
                ),
            })
            step += 1
            expr_h5ad = f"<output_path from step {step - 1}>"

        # Curation step: filter to curated sample list if defined in manifest
        if getattr(manifest, "curated_sample_list", None):
            n_curated = len(manifest.curated_sample_list)
            steps.append({
                "step": step,
                "tool": "dataset_filter_to_curated_samples",
                "key_args": {
                    "adata_path": expr_h5ad,
                    "dataset_id": pid,
                    "out_prefix": f"{pid}_curated",
                },
                "note": (
                    f"Filters to {n_curated} curated samples from the manifest "
                    f"({getattr(manifest, 'curated_sample_source', 'see manifest')}). "
                    "Required before survival or prognostic analysis to remove "
                    "non-target tissue contamination."
                ),
            })
            step += 1
            expr_h5ad = f"<output_path from step {step - 1}>"

    elif source.get("type") == "h5ad":
        # Hosted AnnData .h5ad — single-cell / spatial cohorts (ADR-0006 Role 1,
        # e.g. the Loveless Steele subset). Loaded once per resident MCP process
        # via read_h5ad_cached.
        steps.append({
            "step": step,
            "tool": "decoupler_load_and_visualize_data",
            "key_args": {
                "adata_path": source.get("url"),
                "out_prefix": f"{pid}_loaded",
            },
            "note": (
                "Single-cell / spatial h5ad (ADR-0006 Role 1). Loads the hosted "
                "AnnData via read_h5ad_cached — parsed once per resident MCP "
                "process; the Steele-subset is MB-scale so copy-on-read is fine. "
                "A hosted/private h5ad URL is resolved to a local path via the "
                "authenticated loader (HF_TOKEN). The full integrated atlas is "
                "NOT served on this live path — score its offline-derived "
                "signatures against bulk cohorts with dataset_score_signature."
            ),
        })
        step += 1
        expr_h5ad = f"<output_path from step {step - 1}>"

    else:
        # Unsupported source type — still emit inspect_data with a warning
        expr_h5ad = "<manually loaded h5ad path>"

    steps.append({
        "step": step,
        "tool": "decoupler_inspect_data",
        "key_args": {"adata_path": expr_h5ad},
        "note": (
            f"Expected data_type='{manifest.data_level}', "
            f"analysis_path='{manifest.analysis_path}', "
            f"features_look_like_probes={collapse['required'] and not collapse_precomputed}."
        ),
    })
    step += 1

    if collapse["required"] and not collapse_precomputed:
        steps.append({
            "step": step,
            "tool": "decoupler_collapse_probes_to_genes",
            "key_args": {
                "adata_path": expr_h5ad,
                "gene_symbol_column": collapse.get("gene_symbol_column"),
                "method": collapse["method"],
                "out_prefix": f"{pid}_collapsed",
            },
            "note": "Collapses probe IDs to gene symbols before enrichment tools.",
        })
        de_input = f"<collapsed h5ad from step {step}>"
        step += 1
    else:
        de_input = expr_h5ad

    # PATH P (single-cell / spatial): NOT the bulk DESeq2/limma contrast path.
    # Per-cell work is loaded via read_h5ad_cached; query-time analysis uses the
    # rna_sc per-cell scoring tools, and any group contrast goes through
    # pseudobulk aggregation before the bulk DE / activity fast path.
    if manifest.analysis_path == "P":
        steps.append({
            "step": step,
            "tool": "decoupler_score_transcription_factors",
            "key_args": {"adata_path": de_input},
            "note": (
                "PATH P (single-cell / spatial): score per-cell activity with "
                "the rna_sc tools (decoupler_score_transcription_factors / "
                "_progeny_pathways / _hallmark / _cell_types). For a group "
                "CONTRAST, aggregate to pseudobulk first, then run the bulk DE / "
                "activity fast path — do NOT run DESeq2/limma on per-cell counts."
            ),
        })
        step += 1
        return steps

    # Path A only: filter and preprocess before DESeq2
    if manifest.analysis_path == "A":
        steps.append({
            "step": step,
            "tool": "decoupler_load_and_filter_data",
            "key_args": {
                "counts_path": de_input,
                "group_column": contrast.get("design_factor", "condition"),
                "out_prefix": f"{pid}_filtered",
            },
            "note": "PATH A: Filter low-expression genes before DESeq2.",
        })
        step += 1
        de_input = f"<output_path from step {step - 1}>"

        steps.append({
            "step": step,
            "tool": "decoupler_preprocess_data",
            "key_args": {
                "adata_path": de_input,
                "out_prefix": f"{pid}_preprocessed",
            },
            "note": "PATH A: Normalize, scale, and PCA before DESeq2.",
        })
        step += 1
        de_input = f"<output_path from step {step - 1}>"

    steps.append({
        "step": step,
        "tool": "dataset_validate_contrast",
        "key_args": {
            "metadata_path": de_input,
            "group_column": contrast.get("design_factor", ""),
            "test_group": contrast.get("test_group", ""),
            "control_group": contrast.get("control_group", ""),
        },
        "note": "Confirm groups exist and have enough samples before running DE.",
    })
    step += 1

    de_method = contrast.get("method", "ttest")
    fallback_method = contrast.get("fallback_method")
    steps.append({
        "step": step,
        "tool": "decoupler_differential_expression",
        "key_args": {
            "adata_path": de_input,
            "design_factor": contrast.get("design_factor", "condition"),
            "contrast": [
                contrast.get("design_factor", "condition"),
                contrast.get("test_group", ""),
                contrast.get("control_group", ""),
            ],
            "method": de_method,
            "fallback_to_ttest": fallback_method == "ttest",
            "out_prefix": f"{pid}_de",
        },
        "note": (
            f"analysis_path='{manifest.analysis_path}': "
            + ("use deseq2" if manifest.analysis_path == "A"
               else f"use {de_method}" + (f" (fallback: {fallback_method})" if fallback_method else "") + ", not deseq2")
        ),
    })

    return steps


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------

# ===========================================================================
# Analysis planner
# ===========================================================================

# ---------------------------------------------------------------------------
# Keyword tables (no LLM — pure substring matching on lowercased question)
# ---------------------------------------------------------------------------

_COMPARE_KW: frozenset[str] = frozenset({
    "compare", "comparison", "differ", "difference", "differences",
    "vs", "versus", "between", "higher in", "lower in", "higher than",
    "lower than", "upregulated", "downregulated", "up in", "down in",
    "enriched", "depleted", "differential", "significant", "test vs",
    "control vs",
})

_SCORE_KW: frozenset[str] = frozenset({
    "activity", "activities", "score", "scores", "scoring",
    "which samples", "sample-level", "per sample", "across samples",
    "pathway activity", "tf activity", "hallmark activity",
    "progeny", "collectri", "ranked by", "rank samples",
    "high tf", "high pathway", "low tf", "low pathway",
})

_SURVIVAL_KW: frozenset[str] = frozenset({
    "survival", "prognosis", "prognostic", "overall survival",
    "disease-free", "recurrence", "mortality", "time to event",
    "kaplan", "kaplan-meier", "cox regression", "cox model",
    "hazard ratio", "log-rank", "outcome", "death",
    "associated with survival", "survival analysis",
})

_CORRELATE_KW: frozenset[str] = frozenset({
    "correlate", "correlation", "correlates", "associated with",
    "association", "covariate", "continuous", "regression",
    "trend", "predict", "predictor", "relate to",
    "relation between",
})

_METADATA_KW: frozenset[str] = frozenset({
    "how many", "count", "counts", "sample count", "sample counts",
    "distribution", "group distribution", "subtype distribution",
    "break down", "breakdown", "value counts", "value_counts",
    "how many classical", "how many basal", "how many samples",
    "group sizes", "what groups", "what subtypes", "what categories",
    "list samples", "describe metadata", "summarise metadata",
    "summarize metadata", "metadata summary", "column values",
    "what are the values", "what values",
})

_ALL_KW: dict[str, frozenset[str]] = {
    "metadata_summary": _METADATA_KW,
    "compare_groups": _COMPARE_KW,
    "score_samples": _SCORE_KW,
    "survival": _SURVIVAL_KW,
    "correlate_continuous": _CORRELATE_KW,
}


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _detect_intent(question: str) -> tuple[str, str, list[str]]:
    """
    Return (intent, confidence, matched_keywords).

    Uses substring matching on the lowercased question — no LLM.
    confidence: "high" (≥2 matches), "medium" (1 match), "low" (0 matches).
    Ties broken by insertion order of _ALL_KW (compare_groups has priority).
    """
    q = question.lower()
    scores: dict[str, int] = {
        intent: sum(1 for kw in kws if kw in q)
        for intent, kws in _ALL_KW.items()
    }
    best_intent = max(scores, key=scores.get)
    best_score = scores[best_intent]

    if best_score == 0:
        return "unknown", "low", []

    matched = [kw for kw in _ALL_KW[best_intent] if kw in q]
    confidence = "high" if best_score >= 2 else "medium"
    return best_intent, confidence, matched


def _workflow_for_intent(
    intent: str,
    dataset_id: str,
    manifest: Any | None,
) -> dict[str, Any]:
    """
    Build the ordered workflow steps and metadata for a detected intent.

    All dataset-specific arg values (organism, contrast groups, etc.) are
    drawn from the manifest when available.  Falls back to generic placeholders
    when no manifest is loaded.
    """
    # Extract manifest hints
    analysis_path = "B"
    collapse_required = False
    contrast: dict[str, str] = {}
    organism = "human"

    if manifest is not None:
        analysis_path = getattr(manifest, "analysis_path", "B")
        organism = getattr(manifest, "organism", "human")
        collapse_required = get_collapse_params(manifest)["required"]
        contrast = get_contrast_groups(manifest)

    step_n = 1

    # Preamble: always describe the dataset first
    steps: list[dict] = [{
        "step": step_n,
        "tool": "dataset_describe",
        "purpose": "Confirm data type, analysis path, and default contrast.",
        "status": "available",
        "args_hint": {"dataset_id": dataset_id},
    }]
    step_n += 1

    # Optional loading steps from manifest
    if manifest is not None and hasattr(manifest, "expression_source"):
        if manifest.expression_source.get("type") == "geo_series_matrix":
            steps.append({
                "step": step_n,
                "tool": "decoupler_load_geo_series_matrix",
                "purpose": "Download and parse GEO series matrix to h5ad.",
                "status": "available",
                "args_hint": {
                    "url_or_path": manifest.expression_source.get("url"),
                    "condition_column": (
                        manifest.group_columns[0] if manifest.group_columns else None
                    ),
                },
            })
            step_n += 1

    if collapse_required:
        steps.append({
            "step": step_n,
            "tool": "decoupler_collapse_probes_to_genes",
            "purpose": "Collapse probe IDs to gene symbols (required before enrichment).",
            "status": "available",
            "args_hint": get_collapse_params(manifest) if manifest else {},
        })
        step_n += 1

    # ── Intent-specific steps ────────────────────────────────────────────────
    if intent == "metadata_summary":
        steps += [
            {
                "step": step_n,
                "tool": "dataset_interpret_metadata",
                "purpose": "Show semantic definitions for each metadata column: roles, "
                           "allowed values, missing-value meanings, and analysis rules.",
                "status": "available",
                "args_hint": {
                    "adata_path": "<h5ad path from previous load step>",
                    "dataset_id": dataset_id,
                },
            },
            {
                "step": step_n + 1,
                "tool": "dataset_count_metadata_values",
                "purpose": "Compute value counts dynamically from the loaded data, "
                           "applying manifest missing-value semantics to classify "
                           "annotated vs unannotated samples.",
                "status": "available",
                "args_hint": {
                    "adata_path": "<h5ad path>",
                    "dataset_id": dataset_id,
                },
            },
        ]
        return {
            "steps": steps,
            "required_inputs": ["h5ad path from decoupler_load_geo_series_matrix"],
            "assumptions": [
                "Dataset has been loaded and saved as h5ad.",
                "Counts are computed from the loaded data, not from the manifest.",
            ],
            "warnings": [
                "Sample counts are computed dynamically — they may differ from "
                "documentation if the dataset is subsetted or filtered.",
            ],
            "refusal_conditions": [
                "Refuse to report counts as absolute truth without noting the "
                "data version and any applied subset query.",
            ],
        }

    if intent == "compare_groups":
        de_method = "ttest" if analysis_path == "B" else "deseq2"
        steps += [
            {
                "step": step_n,
                "tool": "dataset_list_valid_sample_groups",
                "purpose": "Confirm available group columns and verify sample counts.",
                "status": "available",
                "args_hint": {"metadata_path": "<h5ad or metadata CSV path>"},
            },
            {
                "step": step_n + 1,
                "tool": "dataset_validate_contrast",
                "purpose": "Check that both groups have enough samples before DE.",
                "status": "available",
                "args_hint": {
                    "group_column": contrast.get("design_factor", "<group column>"),
                    "test_group": contrast.get("test_group", "<test group label>"),
                    "control_group": contrast.get("control_group", "<control group label>"),
                },
            },
            {
                "step": step_n + 2,
                "tool": "decoupler_differential_expression",
                "purpose": "Gene-level differential expression.",
                "status": "available",
                "args_hint": {
                    "method": de_method,
                    "contrast": [
                        contrast.get("design_factor", "condition"),
                        contrast.get("test_group", ""),
                        contrast.get("control_group", ""),
                    ],
                },
            },
            {
                "step": step_n + 3,
                "tool": "decoupler_tf_enrichment_collectri",
                "purpose": "TF activity enrichment from DE statistics.",
                "status": "available",
                "args_hint": {"organism": organism},
            },
            {
                "step": step_n + 4,
                "tool": "decoupler_pathway_enrichment_progeny",
                "purpose": "Pathway enrichment from DE statistics.",
                "status": "available",
                "args_hint": {"organism": organism},
            },
            {
                "step": step_n + 5,
                "tool": "decoupler_hallmark_enrichment",
                "purpose": "Hallmark gene set enrichment from DE statistics.",
                "status": "available",
                "args_hint": {"organism": organism},
            },
        ]
        required_inputs = [
            "Expression file (h5ad, csv, or tsv) — samples as rows, genes as columns",
            "Sample metadata file with group labels",
        ]
        assumptions = [
            f"Expression is pre-normalised (analysis_path='{analysis_path}').",
            "Gene columns are HGNC gene symbols.",
            "At least 3 samples per group.",
        ]
        if analysis_path == "B":
            assumptions.append(
                "Use method='ttest' or 'limma' — NOT 'deseq2' (data is pre-normalised)."
            )
        warnings: list[str] = []
        refusal_conditions: list[str] = []

    elif intent == "score_samples":
        steps += [
            {
                "step": step_n,
                "tool": "dataset_score_bulk_samples",
                "purpose": "Compute per-sample TF/pathway/hallmark activity scores.",
                "status": "available",
                "args_hint": {"resource": "progeny", "organism": organism, "method": "ulm"},
            },
            {
                "step": step_n + 1,
                "tool": "dataset_compare_activity_by_group",
                "purpose": "Optional: compare activity scores between groups.",
                "status": "available",
                "args_hint": {
                    "group_column": contrast.get("design_factor", "<group column>"),
                    "test_group": contrast.get("test_group", ""),
                    "control_group": contrast.get("control_group", ""),
                },
            },
        ]
        required_inputs = [
            "Expression file (h5ad, csv, or tsv) — samples as rows, genes as columns",
        ]
        assumptions = [
            "Expression is pre-normalised.",
            "Gene columns are HGNC gene symbols.",
        ]
        warnings = [
            "Set resource='progeny' for pathway scores, 'collectri' for TF, "
            "'hallmark' for gene sets.",
        ]
        refusal_conditions = []

    elif intent == "survival":
        from src.workflows.survival import check_survival_data_available
        survival_available = (
            check_survival_data_available(manifest)["available"]
            if manifest is not None else False
        )
        survival_status = "available" if survival_available else "not_implemented"
        steps += [
            {
                "step": step_n,
                "tool": "dataset_score_bulk_samples",
                "purpose": "Compute per-sample activity scores for survival modelling.",
                "status": "available",
                "args_hint": {"resource": "progeny", "organism": organism},
            },
            {
                "step": step_n + 1,
                "tool": "survival_analysis",
                "purpose": "Kaplan-Meier or Cox regression with activity as covariate.",
                "status": survival_status,
                "note": (
                    "Survival columns declared in manifest."
                    if survival_available
                    else "Survival data not in manifest. "
                         "Obtain event+time columns from supplementary tables first."
                ),
                "args_hint": {},
            },
        ]
        required_inputs = [
            "Expression file (h5ad, csv, or tsv)",
            "Survival metadata with numeric time and binary event (0/1) columns",
        ]
        assumptions = ["Expression is pre-normalised."]
        warnings = [
            "Survival analysis is NOT YET IMPLEMENTED in this pipeline. "
            "Export activity scores (dataset_score_bulk_samples) and use "
            "lifelines or R survminer externally.",
        ]
        refusal_conditions = [
            "Cannot run survival analysis until event and time columns are confirmed "
            "in the metadata.",
        ]

    elif intent == "correlate_continuous":
        steps += [
            {
                "step": step_n,
                "tool": "dataset_score_bulk_samples",
                "purpose": "Compute per-sample activity scores as correlation input.",
                "status": "available",
                "args_hint": {"resource": "progeny", "organism": organism},
            },
            {
                "step": step_n + 1,
                "tool": "associate_activity_with_covariate",
                "purpose": "Correlate activity with a continuous covariate.",
                "status": "not_implemented",
                "note": (
                    "Not yet implemented. Export activity CSV and use "
                    "scipy.stats.spearmanr or R cor.test externally."
                ),
                "args_hint": {},
            },
        ]
        required_inputs = [
            "Expression file (h5ad, csv, or tsv)",
            "Metadata file with the continuous covariate column",
        ]
        assumptions = [
            "Expression is pre-normalised.",
            "The covariate is a continuous numeric variable.",
        ]
        warnings = [
            "Continuous covariate association is NOT YET IMPLEMENTED. "
            "Export activity scores and correlate externally.",
        ]
        refusal_conditions = []

    else:  # unknown
        steps = [
            {
                "step": 1,
                "tool": "dataset_list_available",
                "purpose": "Discover what datasets are registered.",
                "status": "available",
                "args_hint": {},
            },
            {
                "step": 2,
                "tool": "dataset_describe",
                "purpose": "Understand the dataset before choosing an analysis.",
                "status": "available",
                "args_hint": {"dataset_id": dataset_id},
            },
        ]
        required_inputs = ["Clarification of the analysis goal."]
        assumptions = []
        warnings = [
            "Could not detect a specific analysis intent. "
            "Rephrase to include keywords like: 'compare', 'activity', "
            "'survival', or 'correlate'.",
        ]
        refusal_conditions = [
            "Will not start analysis until the user's goal is clear.",
        ]

    return {
        "steps": steps,
        "required_inputs": required_inputs,
        "assumptions": assumptions,
        "warnings": warnings,
        "refusal_conditions": refusal_conditions,
    }


# ---------------------------------------------------------------------------
# Tool
# ---------------------------------------------------------------------------

def _col_semantics(col_def, col_name: str) -> dict:
    """Counting semantics for a metadata column def (MetadataColumnDef or dict)."""
    if hasattr(col_def, "missing_values"):
        return {
            "missing_set": set(col_def.missing_values),
            "allowed_set": set(col_def.allowed_values),
            "bio_ok": col_def.biological_grouping_allowed,
            "check_col": col_def.decoded_column or col_name,
            "interp_warn": col_def.interpretation_warning,
            "refusal_rules": list(col_def.refusal_rules or []),
        }
    return {
        "missing_set": set(col_def.get("missing_values") or []),
        "allowed_set": set(col_def.get("allowed_values") or []),
        "bio_ok": col_def.get("biological_grouping_allowed", True),
        "check_col": col_def.get("decoded_column") or col_name,
        "interp_warn": col_def.get("interpretation_warning", ""),
        "refusal_rules": list(col_def.get("refusal_rules") or []),
    }


def _classify_metadata_values(raw_counts, missing_set: set, allowed_set: set):
    """Split a value->count mapping into (annotated, missing, unexpected) dicts."""
    annotated: dict[str, int] = {}
    missing: dict[str, int] = {}
    unexpected: dict[str, int] = {}
    for val, cnt in raw_counts.items():
        if val in missing_set or val in {"nan", "None", ""}:
            missing[val] = cnt
        elif allowed_set and val not in allowed_set:
            unexpected[val] = cnt
        else:
            annotated[val] = cnt
    return annotated, missing, unexpected


def _flatten_single_column_summary(
    s: dict, col_def, manifest, column: str, adata_path: str, dataset_id: str
) -> dict:
    """Flattened top-level fields when a single column was requested."""
    empty_meaning = (
        col_def.empty_value_meaning if hasattr(col_def, "empty_value_meaning")
        else col_def.get("empty_value_meaning", "")
    )
    interp_warn = (
        col_def.interpretation_warning if hasattr(col_def, "interpretation_warning")
        else col_def.get("interpretation_warning", "")
    )
    warnings: list[str] = []
    if empty_meaning:
        warnings.append(empty_meaning.strip())
    if interp_warn:
        warnings.append(interp_warn.strip())

    # missing_counts: replace raw empty-string key with the human meaning
    missing_counts: dict[str, int] = {}
    for val, cnt in s["missing"].items():
        label = empty_meaning.split(".")[0].strip() if empty_meaning and val == "" else val
        missing_counts[label or val] = cnt

    # Suggest crosstab when there are missing values.
    crosstab_next: str | None = None
    if s["total_missing"] > 0:
        group_cols = getattr(manifest, "group_columns", [])
        other_cols = [c for c in group_cols if c != column]
        crosstab_col = other_cols[0] if other_cols else "cell_line/tissue"
        crosstab_next = (
            f"To understand what the {s['total_missing']} unannotated samples are, call: "
            f"dataset_crosstab_metadata_values("
            f"adata_path='{adata_path}', "
            f"row_column='{column}', "
            f"col_column='{crosstab_col}', "
            f"dataset_id='{dataset_id}')"
        )

    return {
        "column": s["column"],
        "n_total": s["total"],
        "n_annotated": s["total_annotated"],
        "n_missing": s["total_missing"],
        "annotated_counts": s["annotated"],
        "missing_counts": missing_counts,
        "warnings": warnings,
        "prohibited_inferences": s.get("prohibited_inferences", []),
        "next_step_if_missing": crosstab_next,
    }