agent/nodes.py

"""LangGraph node functions for MOF screening pipeline."""
import json
import re
from pathlib import Path

from agent.state import ScreeningState


TOXICITY_KEYS = (
    "LC50_Pimephales",
    "LC50_Daphnia",
    "IGC50_Tetrahymena",
    "IBC50_Vibrio",
)


def _append_trace(state: ScreeningState, agent: str, action: str, output: dict) -> list[dict]:
    trace = list(state.get("agent_trace", []))
    trace.append({
        "agent": agent,
        "action": action,
        "output": output,
    })
    return trace


def _provider_llm(state: ScreeningState, max_tokens: int = 600):
    provider = state.get("llm_provider", "rule_based")
    api_key = state.get("llm_api_key")

    if provider == "rule_based" or not api_key:
        return None

    from langchain_openai import ChatOpenAI

    if provider == "openai":
        return ChatOpenAI(model="gpt-4o-mini", temperature=0.1, max_tokens=max_tokens, api_key=api_key)
    if provider == "deepseek":
        return ChatOpenAI(
            model="deepseek-chat",
            temperature=0.1,
            max_tokens=max_tokens,
            base_url="https://api.deepseek.com/v1",
            api_key=api_key,
        )
    if provider == "qwen":
        return ChatOpenAI(
            model="qwen-turbo",
            temperature=0.1,
            max_tokens=max_tokens,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
            api_key=api_key,
        )
    return None


def _extract_json_object(text: str) -> dict | None:
    try:
        return json.loads(text)
    except Exception:
        pass

    match = re.search(r"\{.*\}", text, flags=re.S)
    if not match:
        return None
    try:
        return json.loads(match.group(0))
    except Exception:
        return None


def _try_llm_json_agent(state: ScreeningState, agent_name: str, schema_hint: dict) -> dict | None:
    try:
        from agent.prompts import build_agent_json_prompt

        llm = _provider_llm(state)
        if llm is None:
            return None

        response = llm.invoke(build_agent_json_prompt(agent_name, state, schema_hint))
        return _extract_json_object(response.content)
    except Exception:
        return None


def ingest_cif(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))
    cif_path = state["cif_path"]

    try:
        p = Path(cif_path)
        if not p.exists():
            errors.append(f"CIF file not found: {cif_path}")
            return {"errors": errors, "warnings": warnings}
        mof_id = p.stem
    except Exception as e:
        errors.append(f"ingest_cif failed: {e}")
        mof_id = "unknown"

    return {"mof_id": mof_id, "warnings": warnings, "errors": errors}


def six_step_screening_agent(state: ScreeningState) -> dict:
    """Run online screening tools for an uploaded CIF."""
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))
    if errors:
        return {"warnings": warnings, "errors": errors}

    try:
        from tools.online_screening import ACTIVE_TOOLS, run_online_tools

        result = run_online_tools(
            state.get("cif_path"),
            list(ACTIVE_TOOLS),
        )
        return {
            **result,
            "warnings": warnings + list(result.get("warnings", [])),
            "errors": errors + list(result.get("errors", [])),
            "agent_trace": _append_trace(
                state,
                "Online Screening Agent",
                "ran online CIF screening tools",
                {
                    "gate_status": result.get("gate_status"),
                    "recommendation": result.get("recommendation"),
                },
            ) + list(result.get("agent_trace", [])),
        }
    except Exception as e:
        errors.append(f"online_screening failed: {e}")
        return {
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(
                state,
                "Online Screening Agent",
                "online screening failed",
                {"error": str(e)},
            ),
        }


def planning_agent(state: ScreeningState) -> dict:
    """Create a task card before deterministic tools are executed."""
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))
    mof_id = state.get("mof_id", "unknown")

    task_card = {
        "objective": "screen MOF for aromatic VOC adsorption and safe-by-design suitability",
        "mof_id": mof_id,
        "required_tools": [
            "compute_descriptor",
            "predict_adsorption",
            "extract_linker",
            "predict_toxicity",
            "apply_safety_rules",
        ],
        "required_evidence": [
            "SCM descriptor validity",
            "benzene and toluene adsorption predictions",
            "linker identity and extraction confidence",
            "aquatic toxicity endpoints",
            "metal safety tier",
            "PMT and PFAS pre-filter flags",
        ],
        "decision_policy": {
            "adsorption_weight": 0.4,
            "safety_weight": 0.3,
            "toxicity_weight": 0.3,
            "safety_veto": True,
            "audit_penalty_enabled": True,
        },
    }
    execution_plan = {
        "steps": task_card["required_tools"],
        "evidence_mode": "ledger",
        "final_review": [
            "evidence_audit_agent",
            "safety_review_agent",
            "reviewer_panel_agent",
            "decision_agent",
            "report_agent",
        ],
    }

    return {
        "task_card": task_card,
        "execution_plan": execution_plan,
        "warnings": warnings,
        "errors": errors,
        "agent_trace": _append_trace(state, "Planning Agent", "created task card", task_card),
    }


def tool_execution_agent(state: ScreeningState) -> dict:
    """Translate the task card into an executable tool plan."""
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))
    task_card = state.get("task_card") or {}
    required_tools = task_card.get("required_tools", [])
    tool_execution_plan = {
        "mode": "deterministic_tools_with_agent_supervision",
        "tool_order": required_tools,
        "logging_policy": "each tool must write a compact trace entry and evidence ledger record",
        "failure_policy": "continue where possible, then route evidence gaps to audit and repair agents",
    }

    return {
        "tool_execution_plan": tool_execution_plan,
        "warnings": warnings,
        "errors": errors,
        "agent_trace": _append_trace(
            state,
            "Tool Execution Agent",
            "converted task card into ordered tool calls",
            tool_execution_plan,
        ),
    }


def compute_descriptor(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))

    if state.get("errors"):
        return {"warnings": warnings, "errors": errors}

    try:
        from tools.descriptors import compute_scm_eigenvalues
        import yaml

        config_path = Path(__file__).resolve().parent.parent / "configs" / "paths.yaml"
        with open(config_path, "r", encoding="utf-8") as f:
            cfg = yaml.safe_load(f)

        benzene_dim = cfg["adsorption_models"]["benzene_target_dim"]
        toluene_dim = cfg["adsorption_models"]["toluene_target_dim"]

        result_b = compute_scm_eigenvalues(state["cif_path"], benzene_dim)
        result_t = compute_scm_eigenvalues(state["cif_path"], toluene_dim)

        scm_meta = {
            "benzene_eigenvalues": result_b["eigenvalues"],
            "toluene_eigenvalues": result_t["eigenvalues"],
            "raw_dim": result_b["raw_dim"],
            "benzene_padded": result_b["padded"],
            "benzene_truncated": result_b["truncated"],
            "toluene_padded": result_t["padded"],
            "toluene_truncated": result_t["truncated"],
        }

        for r in [result_b, result_t]:
            if r["applicability_warning"]:
                warnings.append(r["applicability_warning"])

        trace_output = {
            "status": "success",
            "raw_dim": scm_meta["raw_dim"],
            "benzene_padded": scm_meta["benzene_padded"],
            "benzene_truncated": scm_meta["benzene_truncated"],
            "toluene_padded": scm_meta["toluene_padded"],
            "toluene_truncated": scm_meta["toluene_truncated"],
        }
        return {
            "scm_meta": scm_meta,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "ran compute_descriptor", trace_output),
        }
    except Exception as e:
        errors.append(f"compute_descriptor failed: {e}")
        trace_output = {"status": "failed", "error": str(e)}
        return {
            "scm_meta": None,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "compute_descriptor failed", trace_output),
        }


def predict_adsorption(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))

    if state.get("scm_meta") is None:
        errors.append("predict_adsorption skipped: no SCM descriptors available.")
        return {"adsorption": None, "warnings": warnings, "errors": errors}

    try:
        from tools.adsorption import predict_benzene, predict_toluene

        meta = state["scm_meta"]
        b = predict_benzene(meta["benzene_eigenvalues"])
        t = predict_toluene(meta["toluene_eigenvalues"])

        adsorption = {
            "benzene_uptake_mg_g": b["uptake_mg_g"],
            "benzene_model": b["model_version"],
            "toluene_uptake_mg_g": t["uptake_mg_g"],
            "toluene_model": t["model_version"],
        }

        for r in [b, t]:
            if r["applicability_warning"]:
                warnings.append(r["applicability_warning"])

        return {
            "adsorption": adsorption,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "ran predict_adsorption", adsorption),
        }
    except Exception as e:
        errors.append(f"predict_adsorption failed: {e}")
        trace_output = {"status": "failed", "error": str(e)}
        return {
            "adsorption": None,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "predict_adsorption failed", trace_output),
        }


def extract_linker(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))

    try:
        from tools.linker_extraction import extract_linker as _extract

        result = _extract(state["cif_path"])
        linker = {
            "metals": result["metals"],
            "linker_smiles": result["linker_smiles"],
            "linker_name": result["linker_name"],
            "linker_formula": result["linker_formula"],
            "extraction_level": result["extraction_level"],
        }
        if result["extraction_level"] == 3:
            warnings.append(result["extraction_note"])

        return {
            "linker": linker,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "ran extract_linker", linker),
        }
    except Exception as e:
        errors.append(f"extract_linker failed: {e}")
        trace_output = {"status": "failed", "error": str(e)}
        return {
            "linker": None,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "extract_linker failed", trace_output),
        }


def predict_toxicity(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))

    linker = state.get("linker")
    smiles = linker.get("linker_smiles") if linker else None

    try:
        from tools.toxicity import predict_toxicity as _predict

        result = _predict(smiles)
        trace_output = {
            key: result.get(key)
            for key in TOXICITY_KEYS
        }
        trace_output["applicability_domain_flag"] = result.get("applicability_domain_flag")
        return {
            "toxicity": result,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "ran predict_toxicity", trace_output),
        }
    except Exception as e:
        errors.append(f"predict_toxicity failed: {e}")
        trace_output = {"status": "failed", "error": str(e)}
        return {
            "toxicity": None,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "predict_toxicity failed", trace_output),
        }


def apply_safety_rules(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))

    try:
        from tools.safety_rules import check_metal_safety, check_pmt_pre_filter

        linker = state.get("linker")
        metals = linker.get("metals", []) if linker else []
        smiles = linker.get("linker_smiles") if linker else None

        metal_result = check_metal_safety(metals) if metals else {
            "tier": "unknown", "flagged_metals": [], "details": "No metals found."
        }
        pmt_result = check_pmt_pre_filter(smiles)

        safety = {
            "metal_tier": metal_result["tier"],
            "metal_flagged": metal_result["flagged_metals"],
            "metal_details": metal_result["details"],
            "pmt_pass": pmt_result["pmt_pass"],
            "pmt_flags": pmt_result["flags"],
            "pmt_descriptors": pmt_result["descriptors"],
        }

        return {
            "safety": safety,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "ran apply_safety_rules", safety),
        }
    except Exception as e:
        errors.append(f"apply_safety_rules failed: {e}")
        trace_output = {"status": "failed", "error": str(e)}
        return {
            "safety": None,
            "warnings": warnings,
            "errors": errors,
            "agent_trace": _append_trace(state, "Tool Execution Agent", "apply_safety_rules failed", trace_output),
        }


def evidence_audit_agent(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))
    ledger = _build_evidence_ledger(state)

    issues = []
    scm_meta = state.get("scm_meta") or {}
    if scm_meta.get("benzene_padded") or scm_meta.get("toluene_padded"):
        issues.append({
            "severity": "medium",
            "source": "compute_descriptor",
            "message": "SCM eigenvalue vector was padded, indicating descriptor dimension mismatch.",
        })
    if scm_meta.get("benzene_truncated") or scm_meta.get("toluene_truncated"):
        issues.append({
            "severity": "medium",
            "source": "compute_descriptor",
            "message": "SCM eigenvalue vector was truncated, indicating descriptor dimension mismatch.",
        })

    linker = state.get("linker") or {}
    if not linker.get("linker_smiles"):
        issues.append({
            "severity": "high",
            "source": "extract_linker",
            "message": "No linker SMILES was identified, so linker toxicity evidence is incomplete.",
        })
    elif linker.get("extraction_level") == 3:
        issues.append({
            "severity": "medium",
            "source": "extract_linker",
            "message": "Linker was identified by degraded matching rather than a high-confidence match.",
        })

    toxicity = state.get("toxicity") or {}
    missing_tox = [k for k in TOXICITY_KEYS if toxicity.get(k) is None]
    if missing_tox:
        issues.append({
            "severity": "high" if len(missing_tox) >= 2 else "medium",
            "source": "predict_toxicity",
            "message": f"Missing toxicity endpoint(s): {', '.join(missing_tox)}.",
        })

    safety = state.get("safety") or {}
    if safety.get("metal_tier") in {"black", "unknown"}:
        issues.append({
            "severity": "high",
            "source": "apply_safety_rules",
            "message": f"Metal tier is {safety.get('metal_tier')}.",
        })
    if not safety.get("pmt_pass", True):
        issues.append({
            "severity": "high",
            "source": "apply_safety_rules",
            "message": "PMT pre-filter failed.",
        })

    if errors:
        issues.append({
            "severity": "high",
            "source": "pipeline",
            "message": "One or more upstream pipeline errors occurred.",
        })

    high_count = sum(1 for issue in issues if issue["severity"] == "high")
    medium_count = sum(1 for issue in issues if issue["severity"] == "medium")
    confidence_penalty = round(min(0.4, high_count * 0.15 + medium_count * 0.07), 2)
    audit_status = "pass" if not issues else "caution"
    if high_count >= 2 or errors:
        audit_status = "fail"

    fallback = {
        "audit_status": audit_status,
        "issues": issues,
        "confidence_penalty": confidence_penalty,
        "requires_repair": audit_status != "pass",
    }
    llm_result = _try_llm_json_agent(state | {"evidence_ledger": ledger}, "Evidence Audit Agent", fallback)
    audit_report = llm_result if isinstance(llm_result, dict) else fallback

    return {
        "evidence_ledger": ledger,
        "audit_report": audit_report,
        "repair_actions": _repair_actions_from_audit(audit_report),
        "warnings": warnings,
        "errors": errors,
        "agent_trace": _append_trace(state, "Evidence Audit Agent", "audited evidence ledger", audit_report),
    }


def repair_agent(state: ScreeningState) -> dict:
    """Create an explicit repair plan when audit finds weak or missing evidence."""
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))
    audit_report = state.get("audit_report") or {}
    issues = audit_report.get("issues", [])
    repair_actions = _repair_actions_from_audit(audit_report)

    repair_steps = []
    for issue in issues:
        source = issue.get("source", "unknown")
        severity = issue.get("severity", "unknown")
        message = issue.get("message", "")
        if source == "compute_descriptor":
            proposed = "re-parse CIF, verify structure validity, then recompute SCM descriptors"
            can_auto_apply = False
        elif source == "extract_linker":
            proposed = "request curated linker identity or run manual linker validation"
            can_auto_apply = False
        elif source == "predict_toxicity":
            proposed = "rerun endpoint models after linker repair or flag endpoint for experiment"
            can_auto_apply = False
        elif source == "apply_safety_rules":
            proposed = "review metal tier, PMT flags, and potential leaching concern"
            can_auto_apply = False
        else:
            proposed = "resolve pipeline issue before final ranking"
            can_auto_apply = False
        repair_steps.append({
            "source": source,
            "severity": severity,
            "issue": message,
            "proposed_action": proposed,
            "can_auto_apply": can_auto_apply,
        })

    repair_report = {
        "repair_status": "not_required" if not repair_steps else "manual_or_external_validation_required",
        "repair_steps": repair_steps,
        "auto_repair_attempted": False,
        "decision_impact": "candidate cannot be promoted to class A until high-severity or required evidence issues are resolved"
        if repair_steps else "no repair penalty",
    }

    return {
        "repair_report": repair_report,
        "repair_actions": repair_actions,
        "warnings": warnings,
        "errors": errors,
        "agent_trace": _append_trace(state, "Repair Agent", "converted audit issues into repair actions", repair_report),
    }


def safety_review_agent(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))
    safety = state.get("safety") or {}
    toxicity = state.get("toxicity") or {}
    audit_report = state.get("audit_report") or {}

    dominant_risks = []
    safety_veto = False
    metal_tier = safety.get("metal_tier", "unknown")
    if metal_tier in {"black", "unknown"}:
        dominant_risks.append(f"{metal_tier} metal tier")
        safety_veto = metal_tier == "black"
    if not safety.get("pmt_pass", True):
        dominant_risks.append("PMT pre-filter failure")
        dominant_risks.extend([
            flag for flag in safety.get("pmt_flags", [])
            if "passed" not in str(flag).lower()
        ])

    valid_tox = [toxicity.get(k) for k in TOXICITY_KEYS if toxicity.get(k) is not None]
    if valid_tox and sum(valid_tox) / len(valid_tox) >= 4.0:
        dominant_risks.append("high predicted aquatic toxicity")

    if audit_report.get("audit_status") == "fail":
        dominant_risks.append("failed evidence audit")

    if safety_veto:
        safety_label = "reject"
    elif dominant_risks:
        safety_label = "caution"
    else:
        safety_label = "pass"

    fallback = {
        "safety_label": safety_label,
        "dominant_risks": dominant_risks,
        "safety_veto": safety_veto,
        "required_validation": _required_validation_actions(state, dominant_risks),
    }
    llm_result = _try_llm_json_agent(state, "Safety Review Agent", fallback)
    safety_review = llm_result if isinstance(llm_result, dict) else fallback

    return {
        "safety_review": safety_review,
        "warnings": warnings,
        "errors": errors,
        "agent_trace": _append_trace(state, "Safety Review Agent", "reviewed safety evidence", safety_review),
    }


def reviewer_panel_agent(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))
    adsorption = state.get("adsorption") or {}
    toxicity = state.get("toxicity") or {}
    audit_report = state.get("audit_report") or {}
    safety_review = state.get("safety_review") or {}

    benzene = adsorption.get("benzene_uptake_mg_g", 0) or 0
    performance_score = round(min(10.0, benzene / 1000 * 10), 1)
    safety_score = 8.0
    if safety_review.get("safety_label") == "caution":
        safety_score = 5.5
    elif safety_review.get("safety_label") == "reject":
        safety_score = 2.0

    valid_tox = [toxicity.get(k) for k in TOXICITY_KEYS if toxicity.get(k) is not None]
    if valid_tox:
        toxicity_penalty = max(0.0, (sum(valid_tox) / len(valid_tox) - 3.0) * 0.8)
        safety_score = round(max(0.0, safety_score - toxicity_penalty), 1)

    practicality_score = 7.5
    if audit_report.get("requires_repair"):
        practicality_score -= 1.5
    if (state.get("linker") or {}).get("extraction_level") == 3:
        practicality_score -= 1.0
    practicality_score = round(max(0.0, practicality_score), 1)

    reviewers = [
        {
            "role": "performance_reviewer",
            "score": performance_score,
            "comment": "Scores adsorption potential using predicted benzene uptake.",
        },
        {
            "role": "safety_reviewer",
            "score": safety_score,
            "comment": "Scores metal, PMT, and predicted aquatic toxicity risk.",
        },
        {
            "role": "practicality_reviewer",
            "score": practicality_score,
            "comment": "Scores evidence completeness and experimental follow-up readiness.",
        },
    ]
    panel_score = round(sum(r["score"] for r in reviewers) / len(reviewers), 2)
    spread = max(r["score"] for r in reviewers) - min(r["score"] for r in reviewers)
    agreement = "high" if spread < 2 else "moderate" if spread < 4 else "low"
    fallback = {
        "reviewers": reviewers,
        "panel_score": panel_score,
        "agreement": agreement,
    }
    llm_result = _try_llm_json_agent(state, "Reviewer Panel Agent", fallback)
    panel = llm_result if isinstance(llm_result, dict) else fallback
    reviewer_reports = panel.get("reviewers", reviewers)

    return {
        "reviewer_reports": reviewer_reports,
        "decision_record": {"panel_score": panel.get("panel_score", panel_score), "panel_agreement": panel.get("agreement", agreement)},
        "warnings": warnings,
        "errors": errors,
        "agent_trace": _append_trace(state, "Reviewer Panel Agent", "completed multi-perspective review", panel),
    }


def decision_agent(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))

    score, base_recommendation, explanation = _compute_rule_based_score(state)
    audit_report = state.get("audit_report") or {}
    safety_review = state.get("safety_review") or {}
    panel_record = state.get("decision_record") or {}

    audit_penalty = float(audit_report.get("confidence_penalty", 0) or 0) * 10
    adjusted_score = round(max(0.0, score - audit_penalty), 2)
    panel_score = panel_record.get("panel_score")
    if isinstance(panel_score, (int, float)):
        adjusted_score = round(0.75 * adjusted_score + 0.25 * float(panel_score), 2)

    safety_veto = bool(safety_review.get("safety_veto"))
    requires_repair = bool(audit_report.get("requires_repair"))

    if safety_veto:
        decision_class = "D"
        agentic_recommendation = "reject_by_safety_rule"
        recommendation = "reject"
    elif requires_repair:
        decision_class = "C"
        agentic_recommendation = "repair_evidence_before_ranking"
        recommendation = "borderline"
    elif adjusted_score >= 7.0:
        decision_class = "A"
        agentic_recommendation = "recommend_for_experimental_validation"
        recommendation = "recommend"
    elif adjusted_score >= 5.0:
        decision_class = "B"
        agentic_recommendation = "validate_before_scaleup"
        recommendation = "borderline"
    else:
        decision_class = "D" if base_recommendation == "reject" else "C"
        agentic_recommendation = "reject" if decision_class == "D" else "repair_evidence_before_ranking"
        recommendation = "reject" if decision_class == "D" else "borderline"

    decision_record = {
        **panel_record,
        "base_score": score,
        "audit_penalty": round(audit_penalty, 2),
        "final_score": adjusted_score,
        "decision_class": decision_class,
        "recommendation": recommendation,
        "agentic_recommendation": agentic_recommendation,
        "safety_veto": safety_veto,
        "blocking_issues": _blocking_issues(state),
        "next_actions": _next_actions(state),
        "repair_status": (state.get("repair_report") or {}).get("repair_status"),
    }

    explanation = (
        f"{explanation} Agentic decision class {decision_class}: {agentic_recommendation}. "
        f"Audit penalty {audit_penalty:.2f}; final agentic score {adjusted_score:.2f}."
    )

    return {
        "final_score": adjusted_score,
        "recommendation": recommendation,
        "decision_record": decision_record,
        "explanation": explanation,
        "warnings": warnings,
        "errors": errors,
        "agent_trace": _append_trace(state, "Decision Agent", "fused scores, audit, and safety review", decision_record),
    }


def report_agent(state: ScreeningState) -> dict:
    warnings = list(state.get("warnings", []))
    errors = list(state.get("errors", []))
    explanation = state.get("explanation")

    llm_report = _try_llm_report(state)
    if llm_report:
        explanation = llm_report

    report_payload = {
        "recommendation": state.get("recommendation"),
        "final_score": state.get("final_score"),
        "decision_record": state.get("decision_record"),
    }

    return {
        "explanation": explanation,
        "warnings": warnings,
        "errors": errors,
        "agent_trace": _append_trace(state, "Report Agent", "generated final report", report_payload),
    }


def score_and_explain(state: ScreeningState) -> dict:
    """Backward-compatible single-node entry point for older imports/tests."""
    decision = decision_agent(state)
    merged = {**state, **decision}
    return report_agent(merged)


def _compute_rule_based_score(state: ScreeningState) -> tuple[float, str, str]:
    import yaml
    config_path = Path(__file__).resolve().parent.parent / "configs" / "thresholds.yaml"
    with open(config_path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f)

    weights = cfg["scoring"]["weights"]
    safety_scores = cfg["scoring"]["safety_scores"]
    thresholds = cfg["scoring"]["recommendation"]
    benzene_ref = cfg["scoring"]["adsorption_normalization"]["benzene_ref_mg_g"]

    adsorption = state.get("adsorption") or {}
    safety = state.get("safety") or {}
    toxicity = state.get("toxicity") or {}

    benzene_uptake = adsorption.get("benzene_uptake_mg_g", 0)
    adsorption_score = min(benzene_uptake / benzene_ref, 1.0) * 10

    tier = safety.get("metal_tier", "unknown")
    safety_score = safety_scores.get(tier, safety_scores["unknown"])
    if not safety.get("pmt_pass", True):
        safety_score = max(safety_score - 3, 0)

    tox_values = [
        toxicity.get("LC50_Pimephales"),
        toxicity.get("LC50_Daphnia"),
        toxicity.get("IGC50_Tetrahymena"),
        toxicity.get("IBC50_Vibrio"),
    ]
    valid_tox = [v for v in tox_values if v is not None]
    if valid_tox:
        mean_tox = sum(valid_tox) / len(valid_tox)
        tox_score = max(0, min(10, (5.0 - mean_tox) * 2 + 5))
    else:
        tox_score = 5.0

    final_score = (
        weights["adsorption"] * adsorption_score
        + weights["safety"] * safety_score
        + weights["toxicity"] * tox_score
    )
    final_score = round(min(10.0, max(0.0, final_score)), 2)

    if final_score >= thresholds["recommend_threshold"]:
        recommendation = "recommend"
    elif final_score >= thresholds["borderline_threshold"]:
        recommendation = "borderline"
    else:
        recommendation = "reject"

    mof_id = state.get("mof_id", "unknown")
    linker_data = state.get("linker") or {}
    metals_str = ", ".join(linker_data.get("metals", []))
    tox_summary = (
        f"mean -log toxicity = {sum(valid_tox)/len(valid_tox):.2f}"
        if valid_tox else "toxicity data unavailable"
    )

    explanation = (
        f"MOF {mof_id} achieves {benzene_uptake:.1f} mg/g benzene uptake. "
        f"Metal node(s) [{metals_str}] are in the {tier} tier. "
        f"Predicted aquatic toxicity: {tox_summary}. "
        f"Final score: {final_score:.2f} -> {recommendation}."
    )

    return final_score, recommendation, explanation


def _build_evidence_ledger(state: ScreeningState) -> list[dict]:
    scm_meta = state.get("scm_meta")
    adsorption = state.get("adsorption")
    linker = state.get("linker")
    toxicity = state.get("toxicity")
    safety = state.get("safety")

    return [
        {
            "tool": "compute_descriptor",
            "status": "success" if scm_meta else "failed",
            "outputs": {
                "raw_dim": (scm_meta or {}).get("raw_dim"),
                "benzene_padded": (scm_meta or {}).get("benzene_padded"),
                "benzene_truncated": (scm_meta or {}).get("benzene_truncated"),
                "toluene_padded": (scm_meta or {}).get("toluene_padded"),
                "toluene_truncated": (scm_meta or {}).get("toluene_truncated"),
            },
            "confidence": "dimension_checked" if scm_meta else "missing",
        },
        {
            "tool": "predict_adsorption",
            "status": "success" if adsorption else "failed",
            "outputs": adsorption or {},
            "confidence": "trained_model" if adsorption else "missing",
        },
        {
            "tool": "extract_linker",
            "status": "success" if linker else "failed",
            "outputs": linker or {},
            "confidence": _linker_confidence(linker),
        },
        {
            "tool": "predict_toxicity",
            "status": "success" if toxicity else "failed",
            "outputs": toxicity or {},
            "confidence": "trained_endpoint_models" if toxicity else "missing",
        },
        {
            "tool": "apply_safety_rules",
            "status": "success" if safety else "failed",
            "outputs": safety or {},
            "confidence": "rule_based" if safety else "missing",
        },
    ]


def _linker_confidence(linker: dict | None) -> str:
    if not linker:
        return "missing"
    level = linker.get("extraction_level")
    if level == 1:
        return "high"
    if level == 2:
        return "medium"
    if level == 3:
        return "low"
    return "unknown"


def _repair_actions_from_audit(audit_report: dict) -> list[str]:
    actions = []
    for issue in audit_report.get("issues", []):
        source = issue.get("source", "")
        if source == "compute_descriptor":
            actions.append("verify CIF parsing and descriptor applicability domain")
        elif source == "extract_linker":
            actions.append("manually validate linker identity or provide curated linker SMILES")
        elif source == "predict_toxicity":
            actions.append("rerun or experimentally validate missing toxicity endpoints")
        elif source == "apply_safety_rules":
            actions.append("review metal tier, PMT flags, and leaching risk")
        elif source == "pipeline":
            actions.append("resolve upstream pipeline errors before ranking")
    return list(dict.fromkeys(actions))


def _required_validation_actions(state: ScreeningState, dominant_risks: list[str]) -> list[str]:
    actions = []
    if any("metal" in risk for risk in dominant_risks):
        actions.append("metal leaching and stability validation")
    if any("PMT" in risk for risk in dominant_risks):
        actions.append("persistence, mobility, and transformation assessment")
    if any("toxicity" in risk for risk in dominant_risks):
        actions.append("experimental aquatic toxicity validation")
    if state.get("audit_report", {}).get("requires_repair"):
        actions.extend(_repair_actions_from_audit(state.get("audit_report", {})))
    return list(dict.fromkeys(actions)) or ["standard adsorption and regeneration validation"]


def _blocking_issues(state: ScreeningState) -> list[str]:
    issues = []
    safety_review = state.get("safety_review") or {}
    audit_report = state.get("audit_report") or {}
    if safety_review.get("safety_veto"):
        issues.append("safety veto triggered")
    for issue in audit_report.get("issues", []):
        if issue.get("severity") == "high":
            issues.append(issue.get("message", "high-severity evidence issue"))
    return issues


def _next_actions(state: ScreeningState) -> list[str]:
    actions = []
    safety_review = state.get("safety_review") or {}
    actions.extend(safety_review.get("required_validation", []))
    actions.extend(state.get("repair_actions", []))
    if not actions:
        actions.append("prioritize for experimental adsorption validation")
    return list(dict.fromkeys(actions))


def _try_llm_report(state: ScreeningState) -> str | None:
    try:
        from agent.prompts import build_report_prompt

        llm = _provider_llm(state, max_tokens=800)
        if llm is None:
            return None
        response = llm.invoke(build_report_prompt(state))
        return response.content
    except Exception:
        return None


def _try_llm_explanation(
    state: ScreeningState, score: float, recommendation: str
) -> str | None:
    provider = state.get("llm_provider", "rule_based")
    api_key = state.get("llm_api_key")

    if provider == "rule_based" or not api_key:
        return None

    try:
        from agent.prompts import build_explanation_prompt
        from langchain_openai import ChatOpenAI

        prompt = build_explanation_prompt(state, score, recommendation)

        if provider == "openai":
            llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3, max_tokens=300, api_key=api_key)
        elif provider == "deepseek":
            llm = ChatOpenAI(
                model="deepseek-chat",
                temperature=0.3,
                max_tokens=300,
                base_url="https://api.deepseek.com/v1",
                api_key=api_key,
            )
        elif provider == "qwen":
            llm = ChatOpenAI(
                model="qwen-turbo",
                temperature=0.3,
                max_tokens=300,
                base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                api_key=api_key,
            )
        else:
            return None

        response = llm.invoke(prompt)
        return response.content
    except Exception:
        return None