tests/test_agent_tools.py

"""Tests for dee/core/agent_tools.py — the tools the Turing agent can call.

Boundary: these tests exercise agent_tools.py's own argument validation and
result shaping, not the underlying dee.core.crispr / dee.core.primers /
dee.core.scoring engines (those have their own test suites) — so
find_guides/design_primers/scoring/evolve are monkeypatched rather than
run for real.
"""

from types import SimpleNamespace

import pandas as pd
import pytest

from dee.core import agent_tools, scoring


# ─── execute_tool dispatch ──────────────────────────────────────────────

def test_execute_tool_unknown_name_returns_error():
    result = agent_tools.execute_tool("not_a_real_tool", {}, False)
    assert result["ok"] is False
    assert "unknown tool" in result["error"]


@pytest.mark.parametrize("name", ["fetch_sequence", "design_crispr_guides", "design_primers"])
def test_execute_tool_requires_signin_for_anonymous(name):
    result = agent_tools.execute_tool(name, {}, True)
    assert result == {
        "ok": False,
        "kind": "signin_required",
        "error": "This requires a free account. Sign in to continue.",
    }


def test_execute_tool_never_raises_on_unexpected_tool_exception(monkeypatch):
    def boom(args):
        raise RuntimeError("unexpected")

    monkeypatch.setitem(agent_tools._TOOLS, "design_crispr_guides",
                         {"fn": boom, "requires_signin": False})
    result = agent_tools.execute_tool("design_crispr_guides", {}, False)
    assert result["ok"] is False
    assert "failed" in result["error"]


# ─── fetch_sequence ──────────────────────────────────────────────────────
# Added 2026-07-11 — chat previously had no way to resolve a gene name /
# accession to a sequence, so a request like "find CRISPR guides in human
# GFP" was a dead end even though dee.core.resolve.resolve_target() (the
# same "paste anything" resolver behind POST /api/crispr/resolve) already
# did exactly this.

def test_fetch_sequence_tool_missing_text_errors():
    result = agent_tools.execute_tool("fetch_sequence", {}, False)
    assert result == {"ok": False, "error": "missing 'text'"}


def test_fetch_sequence_tool_rejects_oversized_text():
    result = agent_tools.execute_tool(
        "fetch_sequence", {"text": "A" * 1_000_001}, False,
    )
    assert result["ok"] is False
    assert "1 Mbp" in result["error"]


def test_fetch_sequence_tool_shapes_successful_result(monkeypatch):
    seen = {}

    def fake_resolve_target(text, organism=""):
        seen.update(text=text, organism=organism)
        return {"ok": True, "kind": "gene", "sequence": "ATGAAA",
                "gene_symbol": "GFP", "label": "GFP · ENST00000 · CDS 6 nt",
                "source": "ensembl"}

    monkeypatch.setattr("dee.core.resolve.resolve_target", fake_resolve_target)
    result = agent_tools.execute_tool(
        "fetch_sequence", {"text": "GFP", "organism": "human"}, False,
    )
    assert seen == {"text": "GFP", "organism": "human"}
    assert result == {
        "ok": True, "kind": "gene", "sequence": "ATGAAA", "length": 6,
        "gene_symbol": "GFP", "label": "GFP · ENST00000 · CDS 6 nt",
        "source": "ensembl",
    }


def test_fetch_sequence_tool_invalid_organism_falls_back_to_empty(monkeypatch):
    seen = {}
    monkeypatch.setattr("dee.core.resolve.resolve_target",
                         lambda text, organism="": seen.update(organism=organism) or {"ok": True, "sequence": "ATG"})
    agent_tools.execute_tool("fetch_sequence", {"text": "GFP", "organism": "mars"}, False)
    assert seen["organism"] == ""


def test_fetch_sequence_tool_passes_through_resolve_failure(monkeypatch):
    monkeypatch.setattr(
        "dee.core.resolve.resolve_target",
        lambda text, organism="": {"ok": False, "error": "Couldn't find “XYZ” in human."},
    )
    result = agent_tools.execute_tool("fetch_sequence", {"text": "XYZ", "organism": "human"}, False)
    assert result == {"ok": False, "error": "Couldn't find “XYZ” in human."}


def test_fetch_sequence_tool_engine_exception_becomes_error_result(monkeypatch):
    def boom(text, organism=""):
        raise RuntimeError("Ensembl is down")

    monkeypatch.setattr("dee.core.resolve.resolve_target", boom)
    result = agent_tools.execute_tool("fetch_sequence", {"text": "GFP", "organism": "human"}, False)
    assert result["ok"] is False
    assert "Ensembl is down" not in result["error"]


# ─── design_crispr_guides ───────────────────────────────────────────────

def _fake_guide(**overrides):
    base = dict(rank=1, strand="+", position=10, spacer="ACGTACGTACGTACGTACGT",
                pam="AGG", composite_score=0.751, on_target_score=0.801,
                gc_pct=50.05, notes="clean")
    base.update(overrides)
    return SimpleNamespace(**base)


def test_crispr_tool_missing_sequence_errors():
    result = agent_tools.execute_tool("design_crispr_guides", {}, False)
    assert result == {"ok": False, "error": "missing 'sequence'"}


def test_crispr_tool_sequence_too_long_errors():
    result = agent_tools.execute_tool(
        "design_crispr_guides", {"sequence": "A" * 1_000_001}, False,
    )
    assert result["ok"] is False
    assert "1 Mbp" in result["error"]


def test_crispr_tool_shapes_and_bounds_args(monkeypatch):
    seen = {}

    def fake_find_guides(sequence, *, enzyme, max_results, mode):
        seen.update(sequence=sequence, enzyme=enzyme, max_results=max_results, mode=mode)
        return [_fake_guide()]

    monkeypatch.setattr("dee.core.crispr.find_guides", fake_find_guides)

    result = agent_tools.execute_tool("design_crispr_guides", {
        "sequence": "ACGT", "enzyme": "WEIRD", "mode": "also weird", "max_results": 9999,
    }, False)

    # Invalid enzyme/mode fall back to the safe defaults rather than erroring
    # (mirrors POST /api/crispr/design's own tolerant handling).
    assert seen["enzyme"] == "cas9"
    assert seen["mode"] == "knockout"
    # Bounded to the chat-sized cap (50), not the REST UI's 500.
    assert seen["max_results"] == 50

    assert result["ok"] is True
    assert result["n_guides"] == 1
    guide = result["guides"][0]
    assert guide["spacer"] == "ACGTACGTACGTACGTACGT"
    assert guide["composite_score"] == 0.751
    assert guide["gc_pct"] == round(50.05, 1)  # rounded to 1 dp


def test_crispr_tool_value_error_from_engine_becomes_error_result(monkeypatch):
    def fake_find_guides(sequence, *, enzyme, max_results, mode):
        raise ValueError("sequence too short for any PAM site")

    monkeypatch.setattr("dee.core.crispr.find_guides", fake_find_guides)
    result = agent_tools.execute_tool("design_crispr_guides", {"sequence": "AC"}, False)
    assert result == {"ok": False, "error": "sequence too short for any PAM site"}


def test_crispr_tool_applies_field_prior_when_populated(monkeypatch):
    # Before 2026-07-11 this tool never consulted the cross-user aggregate at
    # all — it went straight from find_guides to the response, unlike
    # POST /api/crispr/design which always calibrates. Verifies the chat
    # tool now calls the same calibrate_crispr_guides() and reports how many
    # substitution keys were actually blended in.
    monkeypatch.setattr("dee.core.crispr.find_guides", lambda *a, **kw: [_fake_guide()])

    from dee.core import outcomes as O
    fake_prior = O.Prior(tool="crispr", effects={"k1": 0.1, "k2": -0.2}, n_users={}, n_obs={}, weight=0.2)
    monkeypatch.setattr("dee.core.outcomes.load_cached_tool_prior", lambda tool: fake_prior)
    calibrated_with = {}

    def fake_calibrate(guides, prior):
        calibrated_with["guides"] = guides
        calibrated_with["prior"] = prior
        return guides

    monkeypatch.setattr("dee.core.outcomes.calibrate_crispr_guides", fake_calibrate)

    result = agent_tools.execute_tool("design_crispr_guides", {"sequence": "ACGT"}, False)
    assert result["ok"] is True
    assert result["field_prior_keys"] == 2
    assert calibrated_with["prior"] is fake_prior


def test_crispr_tool_field_prior_absent_is_a_noop(monkeypatch):
    # No monkeypatch of load_cached_tool_prior — exercises the real (empty,
    # since no SUPABASE_* env vars in tests) path, confirming the default is
    # silent and harmless rather than an error.
    monkeypatch.setattr("dee.core.crispr.find_guides", lambda *a, **kw: [_fake_guide()])
    result = agent_tools.execute_tool("design_crispr_guides", {"sequence": "ACGT"}, False)
    assert result["ok"] is True
    assert result["field_prior_keys"] == 0


# ─── design_primers ──────────────────────────────────────────────────────

def test_primers_tool_missing_template_errors():
    result = agent_tools.execute_tool("design_primers", {}, False)
    assert result["ok"] is False
    assert "template" in result["error"].lower()


def test_primers_tool_rejects_non_dna_chars():
    result = agent_tools.execute_tool("design_primers", {"template": "ACGTXYZ"}, False)
    assert result["ok"] is False
    assert "A/C/G/T/N" in result["error"]


def test_primers_tool_rejects_oversized_template():
    result = agent_tools.execute_tool("design_primers", {"template": "A" * 70_000}, False)
    assert result["ok"] is False
    assert "bp" in result["error"]


def test_primers_tool_shapes_successful_result(monkeypatch):
    def fake_design_primers(template, target_start, target_end):
        return {
            "ok": True,
            "pairs": [{
                "forward": {"seq": "ACGTACGTACGTACGT", "tm": 59.5, "gc": 50.0, "len": 16},
                "reverse": {"seq": "TTTTAAAACCCCGGGG", "tm": 60.1, "gc": 50.0, "len": 16},
                "product_size": 250,
                "tm_diff": 0.6,
                "cross_dimer": 1,
                "score": 0.9,
            }],
            "n_forward": 4, "n_reverse": 3, "warnings": [],
        }

    monkeypatch.setattr("dee.core.primers.design_primers", fake_design_primers)
    result = agent_tools.execute_tool("design_primers", {"template": "ACGT" * 100}, False)

    assert result["ok"] is True
    assert result["n_forward"] == 4
    assert result["pairs"] == [{
        "forward_seq": "ACGTACGTACGTACGT", "forward_tm": 59.5,
        "reverse_seq": "TTTTAAAACCCCGGGG", "reverse_tm": 60.1,
        "product_size": 250,
    }]
    # Internal candidate metadata (hairpin/self-dimer/score) is dropped —
    # not useful read inline in chat.
    assert "score" not in result["pairs"][0]


def test_primers_tool_passes_through_engine_validation_failure(monkeypatch):
    monkeypatch.setattr("dee.core.primers.design_primers",
                         lambda template, target_start, target_end: {"ok": False, "error": "no pair found"})
    result = agent_tools.execute_tool("design_primers", {"template": "ACGT" * 100}, False)
    assert result == {"ok": False, "error": "no pair found"}


# ─── design_variant_library ──────────────────────────────────────────────

_TEST_PROTEIN = "MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTG"


def _fake_scores_df(n=8):
    return pd.DataFrame({
        "position": list(range(n)), "wt_aa": ["A"] * n, "mut_aa": ["G"] * n,
        "delta_ll": [float(i) for i in range(n)],
    })


def _fake_variant(**overrides):
    base = dict(rank=1, mutation_labels=("A1G", "A2G"), fitness=1.2345)
    base.update(overrides)
    return SimpleNamespace(**base)


def test_variant_library_tool_missing_sequence_errors():
    result = agent_tools.execute_tool("design_variant_library", {}, False)
    assert result == {"ok": False, "error": "missing 'sequence'"}


def test_variant_library_tool_rejects_non_protein_chars():
    result = agent_tools.execute_tool("design_variant_library", {"sequence": "ACGT123"}, False)
    assert result["ok"] is False
    assert "protein" in result["error"].lower()


def test_variant_library_tool_rejects_oversized_sequence():
    result = agent_tools.execute_tool(
        "design_variant_library", {"sequence": "A" * (agent_tools._MAX_PROTEIN_AA + 1)}, False,
    )
    assert result["ok"] is False
    assert str(agent_tools._MAX_PROTEIN_AA) in result["error"]


def test_variant_library_tool_auto_translates_dna_input(monkeypatch):
    # fetch_sequence (and CRISPR/primer tools) all deal in DNA — a model
    # chaining "fetch this gene, then DE it" naturally ends up holding DNA,
    # which used to fail this tool's plain protein-alphabet check with no
    # path forward (found 2026-07-12, looked to the user like the model
    # had forgotten the sequence it had just fetched). ATG + 8x GCC (Ala) +
    # TAA stop = a real, minimal CDS.
    dna = "ATG" + "GCC" * 8 + "TAA"
    monkeypatch.setattr("dee.core.scoring.get_scorer", lambda model: "the-scorer")

    def fake_score_guarded(scorer, protein, *, wait_timeout=None):
        assert protein == "MAAAAAAAA"   # the DNA translated, not passed through raw
        return _fake_scores_df(n=len(protein))

    monkeypatch.setattr("dee.core.scoring.score_guarded", fake_score_guarded)
    monkeypatch.setattr("dee.models.scorer.top_percentile_pool", lambda df, percentile: df)
    monkeypatch.setattr("dee.optimizer.search.evolve", lambda pool, cfg: [_fake_variant(rank=1)])

    result = agent_tools.execute_tool("design_variant_library", {"sequence": dna}, False)
    assert result["ok"] is True
    assert result["n_scored_positions"] == 9


def test_variant_library_tool_dna_translate_failure_falls_back_to_protein_check(monkeypatch):
    # A string that LOOKS like DNA (all A/C/G/T, length divisible by 3) but
    # isn't a valid CDS (no ATG start, so translate_dna's own validation
    # rejects it) must not hard-fail — it falls through to treating the
    # raw text as protein instead. "CCC"*10 is also a valid (if unusual)
    # all-proline protein under the plain amino-acid alphabet check, so it
    # proceeds rather than surfacing a confusing "translation failed" error
    # for something that was never meant to be DNA in the first place.
    not_a_cds = "CCC" * 10   # divisible by 3, all-ACGT, but no ATG start
    monkeypatch.setattr("dee.core.scoring.get_scorer", lambda model: "the-scorer")

    def fake_score_guarded(scorer, protein, *, wait_timeout=None):
        assert protein == not_a_cds   # NOT translated — used as-is
        return _fake_scores_df(n=len(protein))

    monkeypatch.setattr("dee.core.scoring.score_guarded", fake_score_guarded)
    monkeypatch.setattr("dee.models.scorer.top_percentile_pool", lambda df, percentile: df)
    monkeypatch.setattr("dee.optimizer.search.evolve", lambda pool, cfg: [_fake_variant(rank=1)])

    result = agent_tools.execute_tool("design_variant_library", {"sequence": not_a_cds}, False)
    assert result["ok"] is True


def test_variant_library_tool_shapes_and_bounds_args(monkeypatch):
    seen = {}

    monkeypatch.setattr("dee.core.scoring.get_scorer", lambda model: "the-scorer")

    def fake_score_guarded(scorer, protein, *, wait_timeout=None):
        seen.update(scorer=scorer, protein=protein, wait_timeout=wait_timeout)
        return _fake_scores_df()

    monkeypatch.setattr("dee.core.scoring.score_guarded", fake_score_guarded)
    monkeypatch.setattr("dee.models.scorer.top_percentile_pool",
                         lambda df, percentile: df)

    def fake_evolve(pool, cfg):
        seen["k"] = cfg.k
        seen["max_mutations"] = cfg.max_mutations
        seen["min_mutations"] = cfg.min_mutations
        return [_fake_variant(rank=1), _fake_variant(rank=2, mutation_labels=("A3G",), fitness=0.9)]

    monkeypatch.setattr("dee.optimizer.search.evolve", fake_evolve)

    result = agent_tools.execute_tool("design_variant_library", {
        "sequence": _TEST_PROTEIN, "host": "WEIRD", "k": 9999, "max_mutations": 99,
    }, False)

    assert seen["protein"] == _TEST_PROTEIN
    assert seen["wait_timeout"] == 20.0
    # Invalid host falls back to the safe default rather than erroring.
    assert seen["k"] == 20        # clamped to the chat cap, not the REST 200
    assert seen["max_mutations"] == 6
    assert seen["min_mutations"] <= seen["max_mutations"]

    assert result["ok"] is True
    assert result["host"] == "e_coli"
    assert result["n_scored_positions"] == 8
    assert result["variants"] == [
        {"rank": 1, "mutations": ["A1G", "A2G"], "fitness": round(1.2345, 3)},
        {"rank": 2, "mutations": ["A3G"], "fitness": 0.9},
    ]


def test_variant_library_tool_returns_real_dna_per_variant(monkeypatch):
    # Before 2026-07-12 this tool stopped at mutation labels ("W44K") with
    # no actual sequence — nothing to copy/paste or order. Uses REAL
    # search.Mutation/Variant (not the loose SimpleNamespace _fake_variant
    # above, which lacks .mutations and would silently fail inside
    # variants_to_dataframe's apply_variant() — a gap that let the earlier,
    # broken version of this change pass its own tests without actually
    # exercising the encoding path).
    from dee.optimizer.search import Mutation, Variant

    monkeypatch.setattr("dee.core.scoring.get_scorer", lambda model: "the-scorer")
    monkeypatch.setattr("dee.core.scoring.score_guarded",
                         lambda scorer, protein, **kw: _fake_scores_df())
    monkeypatch.setattr("dee.models.scorer.top_percentile_pool", lambda df, percentile: df)

    # _TEST_PROTEIN[0] == "M", _TEST_PROTEIN[1] == "S" — real WT residues,
    # so apply_variant()'s own mismatch check passes.
    real_variant = Variant(
        mutations=(Mutation(position=0, wt_aa="M", mut_aa="A", delta_ll=1.0),
                   Mutation(position=1, wt_aa="S", mut_aa="G", delta_ll=0.5)),
        fitness=1.5, rank=1,
    )
    monkeypatch.setattr("dee.optimizer.search.evolve", lambda pool, cfg: [real_variant])

    result = agent_tools.execute_tool(
        "design_variant_library", {"sequence": _TEST_PROTEIN, "host": "e_coli"}, False,
    )
    assert result["ok"] is True
    variant = result["variants"][0]
    assert variant["mutations"] == ["M1A", "S2G"]
    assert variant["protein"][:2] == "AG"          # M1A, S2G applied
    assert variant["protein"][2:] == _TEST_PROTEIN[2:]   # rest unchanged
    assert isinstance(variant["dna"], str) and len(variant["dna"]) > 0
    assert variant["dna"].strip("ACGT") == ""       # real, plain DNA — no placeholder text
    # AA positions, not an nt-level diff against WT — see the long comment
    # at this field's definition in agent_tools.py for why.
    assert variant["mutated_positions"] == [0, 1]
    assert variant["length_bp"] == len(variant["dna"])


def test_variant_library_tool_busy_returns_kind_busy(monkeypatch):
    monkeypatch.setattr("dee.core.scoring.get_scorer", lambda model: "the-scorer")

    def fake_score_guarded(scorer, protein, *, wait_timeout=None):
        raise scoring.ScoringBusyError("busy")

    monkeypatch.setattr("dee.core.scoring.score_guarded", fake_score_guarded)
    result = agent_tools.execute_tool("design_variant_library", {"sequence": _TEST_PROTEIN}, False)
    assert result["ok"] is False
    assert result["kind"] == "busy"


def test_variant_library_tool_scoring_exception_becomes_error_result(monkeypatch):
    monkeypatch.setattr("dee.core.scoring.get_scorer", lambda model: "the-scorer")

    def fake_score_guarded(scorer, protein, *, wait_timeout=None):
        raise RuntimeError("torch blew up")

    monkeypatch.setattr("dee.core.scoring.score_guarded", fake_score_guarded)
    result = agent_tools.execute_tool("design_variant_library", {"sequence": _TEST_PROTEIN}, False)
    assert result["ok"] is False
    assert "torch blew up" not in result["error"]


def test_variant_library_tool_applies_field_prior_when_populated(monkeypatch):
    # Before 2026-07-11 this tool went straight from scoring to search — it
    # never consulted the cross-user aggregate at all, unlike POST /api/run
    # which always blends it into round-1 scores. Verifies the chat tool now
    # does the same blend and reports how many substitution types applied.
    monkeypatch.setattr("dee.core.scoring.get_scorer", lambda model: "the-scorer")
    monkeypatch.setattr("dee.core.scoring.score_guarded",
                         lambda scorer, protein, **kw: _fake_scores_df())
    monkeypatch.setattr("dee.models.scorer.top_percentile_pool", lambda df, percentile: df)

    from dee.core import aggregate as agg
    fake_prior = agg.GlobalPrior(effects={("A", "G"): 2.0}, n_users={}, n_obs={})
    monkeypatch.setattr("dee.core.aggregate.load_cached_global_prior", lambda: fake_prior)

    seen = {}

    def fake_evolve(pool, cfg):
        seen["delta_ll"] = list(pool["delta_ll"])
        return [_fake_variant(rank=1)]

    monkeypatch.setattr("dee.optimizer.search.evolve", fake_evolve)

    result = agent_tools.execute_tool("design_variant_library", {"sequence": _TEST_PROTEIN}, False)
    assert result["ok"] is True
    assert result["field_prior_substitutions"] == 1
    # Every row is (A -> G) in _fake_scores_df, weight 0.3 * effect 2.0 = +0.6
    # on top of the original delta_ll values (0..7).
    assert seen["delta_ll"] == [float(i) + 0.6 for i in range(8)]


def test_variant_library_tool_field_prior_absent_is_a_noop(monkeypatch):
    # No monkeypatch of load_cached_global_prior — exercises the real
    # (empty, since no SUPABASE_* env vars in tests) path, confirming the
    # default is silent and harmless rather than an error.
    monkeypatch.setattr("dee.core.scoring.get_scorer", lambda model: "the-scorer")
    monkeypatch.setattr("dee.core.scoring.score_guarded",
                         lambda scorer, protein, **kw: _fake_scores_df())
    monkeypatch.setattr("dee.models.scorer.top_percentile_pool", lambda df, percentile: df)
    monkeypatch.setattr("dee.optimizer.search.evolve", lambda pool, cfg: [_fake_variant(rank=1)])

    result = agent_tools.execute_tool("design_variant_library", {"sequence": _TEST_PROTEIN}, False)
    assert result["ok"] is True
    assert result["field_prior_substitutions"] == 0


def test_variant_library_tool_evolve_value_error_becomes_error_result(monkeypatch):
    monkeypatch.setattr("dee.core.scoring.get_scorer", lambda model: "the-scorer")
    monkeypatch.setattr("dee.core.scoring.score_guarded",
                         lambda scorer, protein, **kw: _fake_scores_df())
    monkeypatch.setattr("dee.models.scorer.top_percentile_pool", lambda df, percentile: df)

    def fake_evolve(pool, cfg):
        raise ValueError("Filtered mutation pool is empty")

    monkeypatch.setattr("dee.optimizer.search.evolve", fake_evolve)
    result = agent_tools.execute_tool("design_variant_library", {"sequence": _TEST_PROTEIN}, False)
    assert result == {"ok": False, "error": "Filtered mutation pool is empty"}