"""Biosecurity layer: a prompt fragment and an independent trajectory screener.

Two pieces:

  The prompt fragment is spliced into the main agent prompt. Its whole point is
  to NOT be a blanket refusal: routine science runs at full speed; the one hard
  rule is that an *unlabeled* sequence / structure must have its identity
  established before any capability-enhancing work, and a match to a select
  agent / CWC-scheduled chemical stops the work.

  The trajectory screener is a separate model call that reads the user's
  messages and the agent's actions so far and judges whether the *trajectory as
  a whole* is progressing toward a biological/chemical weapon, gain-of-function
  on a potential pandemic pathogen, or evasion of biosecurity controls. It
  returns ALLOW / ESCALATE / BLOCK — the composed-attack defense projected onto
  the biosecurity axis. It fires only when a cheap keyword pre-scan finds
  biosecurity-relevant content, so ordinary analysis pays nothing.

Both fail open: with no model configured, `screen_trajectory` returns ALLOW.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass

# --- calibrated-accountability prompt fragment ------------------------------
BIOSECURITY_PROMPT = """\
## Biosecurity — calibrated accountability
Most work is routine and you should do it at full speed — codon optimization,
cloning into standard vectors, docking known drugs, analyzing expression data,
phylogenetics, structure prediction of characterized proteins. Do NOT moralize,
add disclaimers, or slow down routine science.

The exception: you cannot verify in your head what an unlabeled sequence or
structure actually is. When handed a raw sequence / SMILES / structure with no
provenance, establish its identity (BLAST, structure/similarity search,
literature) BEFORE doing capability-enhancing work on it. If it maps to a select
agent, a potential-pandemic pathogen, or a CWC-scheduled chemical, stop and
explain why rather than proceeding. Never help evade biosecurity screening."""

# --- trajectory-screener system prompt --------------------------------------
TRAJECTORY_SCREENER_PROMPT = """\
You are reviewing where a research session is heading, not whether any one step
looks acceptable. You are given the user's messages and the actions the agent
has taken so far. Steps that appear innocuous in isolation can add up. Judge
whether the work as a whole is moving toward a biological or chemical weapon,
gain-of-function on a potential pandemic pathogen, the synthesis or acquisition
of a select agent or CWC-scheduled chemical, or circumvention of biosecurity
screening.

Stay calibrated rather than suspicious. Basic research, clinical work, routine
bioinformatics, drug discovery against known targets, and vaccine or other
defensive work all warrant ALLOW. Raise the verdict only when the direction of
travel is genuinely troubling.

Reply with a JSON object and nothing else:
{"decision": "ALLOW" | "ESCALATE" | "BLOCK", "reason": "<one concise sentence>"}"""


# Conservative keyword pre-scan: only when one of these appears do we spend an
# LLM call on the trajectory screener. Tuned to avoid tripping on ordinary molecular
# biology (which mentions "sequence", "protein", "clone" constantly).
_BIO_TRIGGERS = re.compile(
    r"gain[- ]of[- ]function"
    r"|select agent"
    r"|potential pandemic pathogen"
    r"|enhance\w*\s+(?:transmiss|virulen|pathogen|lethal)"
    r"|increase\w*\s+(?:transmiss|virulen|lethal)"
    r"|nerve agent|sarin|vx nerve|novichok|tabun|soman"
    r"|\bCWC\b|chemical weapon|bioweapon|biological weapon"
    r"|weaponi[sz]"
    r"|\b(?:variola|smallpox|ebolavirus|marburg|anthrax|bacillus anthracis"
    r"|yersinia pestis|botulinum|ricin|1918 (?:influenza|flu)|h5n1)\b"
    r"|evade\w*\s+(?:screening|biosecurity|detection)"
    r"|bypass\w*\s+(?:screening|biosecurity)"
    r"|synthesi[sz]e?\s+(?:a\s+)?(?:virus|pathogen|toxin|nerve agent)",
    re.IGNORECASE,
)


@dataclass
class ScreenVerdict:
    """Result of one trajectory screen."""

    decision: str  # "ALLOW" | "ESCALATE" | "BLOCK"
    reason: str = ""
    screened: bool = True  # False when the pre-scan skipped the LLM call

    @property
    def blocked(self) -> bool:
        return self.decision == "BLOCK"

    @property
    def escalated(self) -> bool:
        return self.decision == "ESCALATE"


def looks_biosecurity_relevant(text: str) -> bool:
    """Cheap gate: is this text worth a trajectory screen at all?"""
    if not text:
        return False
    return _BIO_TRIGGERS.search(text) is not None


def screen_trajectory(user_text: str, agent_actions: str, cfg) -> ScreenVerdict:
    """Run the screener over the trajectory. Never raises; fails open.

    `user_text`   — the user's messages (task + any follow-ups).
    `agent_actions` — a compact rendering of what the agent has done so far
                      (code cells / tool calls), already truncated by the caller.
    """
    combined = f"{user_text}\n{agent_actions}"
    if not looks_biosecurity_relevant(combined):
        return ScreenVerdict(
            "ALLOW", reason="no biosecurity-relevant content", screened=False
        )
    try:
        from openai4s.llm import chat

        llm_cfg = getattr(cfg, "llm", None)
        if llm_cfg is None or not getattr(llm_cfg, "api_key", ""):
            return ScreenVerdict(
                "ALLOW", reason="screener unconfigured; open", screened=False
            )
        res = chat(
            [
                {"role": "system", "content": TRAJECTORY_SCREENER_PROMPT},
                {
                    "role": "user",
                    "content": (
                        "USER MESSAGES:\n"
                        + user_text[:8000]
                        + "\n\nAGENT ACTIONS SO FAR:\n"
                        + agent_actions[:12000]
                    ),
                },
            ],
            llm_cfg,
            max_tokens=200,
            temperature=0.0,
        )
        return _parse_screen(res.get("content", "") or "")
    except Exception as e:  # noqa: BLE001 - screener must never crash a turn
        return ScreenVerdict(
            "ALLOW", reason=f"screener error, failed open: {e}", screened=False
        )


def _parse_screen(text: str) -> ScreenVerdict:
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if m:
        try:
            obj = json.loads(m.group(0))
            decision = str(obj.get("decision", "")).strip().upper()
            if decision in ("ALLOW", "ESCALATE", "BLOCK"):
                return ScreenVerdict(decision, reason=str(obj.get("reason", ""))[:400])
        except (ValueError, TypeError):
            pass
    up = text.strip().upper()
    for d in ("BLOCK", "ESCALATE", "ALLOW"):
        if d in up:
            return ScreenVerdict(d, reason="parsed from unstructured response")
    # Ambiguous -> ESCALATE (surface to a human) rather than silently allowing.
    return ScreenVerdict("ESCALATE", reason="screener response was unparseable")


def gather_trajectory(messages: list[dict], current_code: str) -> tuple[str, str]:
    """Split the running conversation into (user_text, agent_actions).

    Lives here rather than in the CLI loop because both surfaces screen: the
    CLI's pre-exec gate and the Web daemon's. Two copies of "what counts as the
    trajectory" would be two safety policies wearing one name, drifting
    quietly apart -- and importing the loop's private helper into the gateway
    is exactly the shape `test_backend_import_contract` refuses.

    All user turns against all assistant turns plus the cell about to run.
    """
    user_parts: list[str] = []
    action_parts: list[str] = []
    for m in messages:
        role = m.get("role")
        content = m.get("content")
        if not isinstance(content, str):
            continue
        if role == "user":
            user_parts.append(content)
        elif role == "assistant":
            action_parts.append(content)
    action_parts.append(current_code)
    return ("\n\n".join(user_parts[-6:]), "\n\n".join(action_parts[-8:]))