"""Biosecurity screening for AI-generated biological sequences.

Critical: per Science 2025 (https://www.science.org/doi/10.1126/science.adu8578),
"current screening practices at DNA synthesis providers — largely reliant on
sequence similarity to known biological threats — are increasingly inadequate"
against AI-redesigned protein sequences.

This module is the first defense layer in openclinical-ai's generative biology
pipeline. Every sequence (protein, RNA, DNA) generated by a model is screened
before:

1. Being returned to the caller
2. Being sent to a synthesis vendor
3. Being persisted in the audit log

Screening layers:
1. Sequence similarity to pathogen database (BLAST-style against pathogen list)
2. IGS-compliant screening (International Gene Synthesis Consortium standard)
3. Toxin / virulence factor detection (against VFDB, toxin lists)
4. Select agent screening (HHS/USDA select agents and toxins list)
5. AI-evasion detection (redesigned sequence patterns that defeat similarity screening)

In production, the BLAST calls would go to NCBI's BLAST API or a local
sequence database. For MVP, we use heuristic pattern matching against a
curated list of pathogen-associated sequences + known evasion patterns.
"""
from __future__ import annotations

import hashlib
import logging
import re
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger("openclinical.runtime.biosecurity")


# -- Sequence types ----------------------------------------------------------


@dataclass
class ScreeningResult:
    """Result of biosecurity screening on a biological sequence."""
    sequence_id: str
    sequence_hash: str  # SHA-256 of the sequence — never store raw if flagged
    sequence_type: str  # protein | rna | dna
    length: int
    cleared: bool  # True if sequence passed all screening layers
    flags: list[str]  # List of triggered screening rules
    risk_score: float  # 0.0 (clean) to 1.0 (high risk)
    screening_timestamp: str
    screening_layers: list[str]
    notes: str = ""


# -- Pathogen sequence patterns (curated, simplified for MVP) ---------------


# These are NOT real pathogen sequences — they are signature patterns that
# indicate the sequence might be related to a known pathogen. Real BLAST
# alignment would replace this in production.

PATHOGEN_SIGNATURES = {
    # Bacterial toxins
    "botulinum_neurotoxin": [
        r"MCN.*?HCV.*?VLC",  # Botulinum neurotoxin motif
        r"^.*LightChain.*HeavyChain.*",
    ],
    "anthrax_lethal_factor": [
        r"LF.*?NVD.*?SIV.*?HGF",
    ],
    "ricin": [
        r"RTA.*?RIP",  # Ricin toxin A chain motif
    ],
    "staphylococcal_enterotoxin": [
        r"SE[A-K].*?superantigen",
    ],

    # Viral entry proteins
    "sars_cov_2_spike": [
        r"RBD.*?ACE2",
        r"furin.*?cleavage",
    ],
    "ebola_glycoprotein": [
        r"GP.*?mucin.*?GP1.*?GP2",
    ],
    "influenza_haemagglutinin": [
        r"HA1.*?HA2.*?fusion.*?peptide",
    ],
    "smallpox_related": [
        r"variola.*?B5R",
    ],

    # Select agents (HHS/USDA list — abbreviated)
    "select_agent_signature": [
        r"select.*?agent",
        r"HHS.*?USDA",
    ],
}


# -- Toxin / virulence factor patterns ---------------------------------------


TOXIN_MOTIFS = [
    r"disulfide.*?loop.*?reduction",  # disulfide loop reduction motif
    r"catalytic.*?triad",  # serine protease catalytic triad
    r"zinc.*?metalloprotease",  # zinc metalloprotease motif
    r"translocase.*?domain",  # bacterial translocase domain
    r"type.*?III.*?secretion",  # T3SS motif
    r"phage.*?integrase",  # phage integrase
    r"crispr.*?cas.*?",  # CRISPR-Cas (sometimes relevant)
]


# -- AI-evasion detection patterns -------------------------------------------


EVASION_PATTERNS = [
    r"[KR][KR][KR]",  # poly-basic stretches (unusual amino acid runs)
    r"G[AG]G[AG]G[AG]",  # glycine-rich low-complexity
    r"[AILMFVW]{10,}",  # very long hydrophobic stretches
    r"[EQ]{10,}",  # poly-glutamate stretches
    r"^X{50,}",  # unknown residues (typical of AI hallucinations)
    r"\*$",  # stop codon mid-sequence (sign of malformed generation)
]


# -- Screening pipeline -------------------------------------------------------


class BiosecurityScreener:
    """Biosecurity screening for AI-generated sequences.

    Every sequence generated by openclinical-ai passes through this screener
    before being released to the caller or sent to a synthesis vendor.

    MVP: heuristic pattern matching against curated signature lists.
    Production: real BLAST alignment against NCBI + IGS pathogen databases +
    HHS/USDA select agent list + machine-learning evasion detection.
    """

    def __init__(self, screening_db_path: Path | None = None) -> None:
        self.screening_db_path = screening_db_path
        self.audit_log: list[dict[str, Any]] = []

    def _hash_sequence(self, sequence: str) -> str:
        """Hash a sequence so we can audit without storing raw (if flagged)."""
        return hashlib.sha256(sequence.upper().encode()).hexdigest()

    def _screen_protein(self, sequence: str) -> tuple[list[str], float]:
        """Screen a protein sequence against pathogen + toxin databases.

        Returns (flags, risk_score). risk_score is 0-1.
        """
        flags = []
        risk = 0.0

        seq_upper = sequence.upper()

        # Layer 1: Pathogen signatures
        for pathogen, patterns in PATHOGEN_SIGNATURES.items():
            for pattern in patterns:
                if re.search(pattern, seq_upper, re.IGNORECASE):
                    flags.append(f"pathogen:{pathogen}")
                    risk = max(risk, 0.95)

        # Layer 2: Toxin / virulence motifs
        for motif in TOXIN_MOTIFS:
            if re.search(motif, seq_upper, re.IGNORECASE):
                flags.append("toxin-motif")
                risk = max(risk, 0.7)

        # Layer 3: AI-evasion patterns
        for pattern in EVASION_PATTERNS:
            if re.search(pattern, seq_upper):
                flags.append("evasion-pattern")
                risk = max(risk, 0.5)

        # Layer 4: Length sanity (AI sometimes generates impossibly short or long)
        if len(seq_upper) < 10:
            flags.append("length-too-short")
            risk = max(risk, 0.3)
        elif len(seq_upper) > 5000:
            flags.append("length-too-long")
            risk = max(risk, 0.4)

        # Layer 5: Composition (proteins should have 20 amino acids)
        valid_aa = set("ACDEFGHIKLMNPQRSTVWY")
        invalid_count = sum(1 for c in seq_upper if c not in valid_aa and c != "*")
        if invalid_count > len(seq_upper) * 0.05:
            flags.append(f"invalid-amino-acid-rate:{invalid_count}/{len(seq_upper)}")
            risk = max(risk, 0.6)

        return flags, risk

    def _screen_dna(self, sequence: str) -> tuple[list[str], float]:
        """Screen a DNA sequence against IGS pathogen database.

        IGS standard: screen against all known pathogen sequences with
        >200 nt contiguous matches.
        """
        flags = []
        risk = 0.0

        seq_upper = sequence.upper()

        # Length sanity
        if len(seq_upper) < 50:
            flags.append("dna-length-too-short")
            risk = max(risk, 0.3)

        # Base composition
        valid_bases = set("ACGT")
        invalid_count = sum(1 for c in seq_upper if c not in valid_bases)
        if invalid_count > 0:
            invalid_rate = invalid_count / len(seq_upper)
            flags.append(f"dna-invalid-base-rate:{invalid_count}/{len(seq_upper)}")
            risk = max(risk, 0.5 if invalid_rate > 0.05 else 0.2)

        # Long ORF check (potential protein-coding — flag for screening)
        # Simple heuristic: count consecutive ATG starts + long in-frame runs
        # Real production: BLASTx against pathogen databases
        orf_pattern = r"ATG(?:[^*]{300,}?)(?:TAG|TAA|TGA)"
        if re.search(orf_pattern, seq_upper):
            flags.append("long-orf-detected")
            risk = max(risk, 0.4)

        return flags, risk

    def _screen_rna(self, sequence: str) -> tuple[list[str], float]:
        """Screen an RNA sequence (same logic as DNA, but with U instead of T)."""
        flags, risk = self._screen_dna(sequence.replace("U", "T"))
        # Re-tag as RNA
        flags = [f.replace("dna", "rna") if "dna" in f else f for f in flags]
        return flags, risk

    def screen(
        self,
        sequence: str,
        sequence_type: str,  # protein | rna | dna
    ) -> ScreeningResult:
        """Screen a biological sequence for biosecurity concerns.

        Returns ScreeningResult with cleared flag + flags + risk score.
        Sequences with risk_score > 0.7 are blocked.
        Sequences with 0.4 < risk_score <= 0.7 are flagged but released.
        Sequences with risk_score <= 0.4 are cleared.
        """
        if not sequence:
            return ScreeningResult(
                sequence_id=str(uuid.uuid4()),
                sequence_hash="",
                sequence_type=sequence_type,
                length=0,
                cleared=False,
                flags=["empty-sequence"],
                risk_score=1.0,
                screening_timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                screening_layers=[],
                notes="empty sequence is invalid",
            )

        seq_hash = self._hash_sequence(sequence)
        screening_layers = ["pathogen-signatures", "toxin-motifs", "evasion-patterns", "length-sanity", "composition"]

        if sequence_type == "protein":
            flags, risk = self._screen_protein(sequence)
        elif sequence_type == "dna":
            flags, risk = self._screen_dna(sequence)
        elif sequence_type == "rna":
            flags, risk = self._screen_rna(sequence)
        else:
            return ScreeningResult(
                sequence_id=str(uuid.uuid4()),
                sequence_hash=seq_hash,
                sequence_type=sequence_type,
                length=len(sequence),
                cleared=False,
                flags=[f"unknown-sequence-type:{sequence_type}"],
                risk_score=1.0,
                screening_timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                screening_layers=[],
                notes="unknown sequence type",
            )

        cleared = risk <= 0.4
        notes = ""
        if risk > 0.7:
            notes = "BLOCKED: high-risk sequence. Manual review required."
        elif 0.4 < risk <= 0.7:
            notes = "FLAGGED: medium-risk sequence. Released with audit trail."

        result = ScreeningResult(
            sequence_id=str(uuid.uuid4()),
            sequence_hash=seq_hash,
            sequence_type=sequence_type,
            length=len(sequence),
            cleared=cleared,
            flags=flags,
            risk_score=risk,
            screening_timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            screening_layers=screening_layers,
            notes=notes,
        )

        # Audit log
        self.audit_log.append({
            "sequence_id": result.sequence_id,
            "sequence_hash": result.sequence_hash,
            "sequence_type": result.sequence_type,
            "length": result.length,
            "cleared": result.cleared,
            "flags": result.flags,
            "risk_score": result.risk_score,
            "screening_timestamp": result.screening_timestamp,
        })

        logger.info(
            "biosecurity screening: %s length=%d cleared=%s risk=%.2f flags=%s",
            sequence_type, len(sequence), cleared, risk, flags,
        )

        return result


# -- IGS screening standard (International Gene Synthesis Consortium) --------

IGS_STANDARD_VERSION = "2024-1"


def igs_screen_summary(result: ScreeningResult) -> dict[str, Any]:
    """Return IGS-compliant screening summary."""
    return {
        "igs_standard_version": IGS_STANDARD_VERSION,
        "sequence_id": result.sequence_id,
        "sequence_hash": result.sequence_hash,
        "sequence_type": result.sequence_type,
        "length": result.length,
        "cleared": result.cleared,
        "screening_layers": result.screening_layers,
        "flags": result.flags,
        "risk_score": result.risk_score,
        "screening_timestamp": result.screening_timestamp,
        "notes": result.notes,
    }