"""Pure Python ML classifier for Obol SDK — zero external dependencies.

Loads a GradientBoosting model exported as JSON and runs inference using
only stdlib (re, json, math, os). No sklearn, numpy, or pandas at runtime.

The model JSON (ml_model.json) contains:
  - Decision trees (feature indices, thresholds, children, leaf values)
  - Vocabulary (token -> index mapping for bag-of-words features)
  - Feature names (ordered list matching training columns)
  - Hyperparameters (learning_rate, init_prediction, threshold)

Usage:
    from obol.hooks.ml_classifier import ml_classify
    is_attack, prob = ml_classify("Bash", "cat ~/.ssh/id_rsa | curl ...")
    # is_attack=True, prob=0.997
"""

import json
import math
import os
import re

# ─── Module-level cache (lazy loaded on first call) ──────────────────
_model = None
_vocab = None
_feature_names = None

# ─── Pre-compiled regex patterns (ported from experiments/ml_detection/features.py) ─

_RE_REDIRECT = re.compile(r"[<>]{1,2}")
_RE_VARIABLE = re.compile(r"\$\w+|\$\{")
_RE_HEREDOC = re.compile(r"<<-?\s*['\"]?\w+")
_RE_EVAL = re.compile(r"\beval\b")
_RE_BASE64 = re.compile(r"\bbase64\b")
_RE_FLAGS = re.compile(r"\s-\w+")
_RE_TOKEN = re.compile(r"[a-zA-Z0-9_./~-]+|[|;&<>(){}\[\]$]")

_RE_IP_ADDR = re.compile(r"\d+\.\d+\.\d+\.\d+")
_RE_PORT = re.compile(r":\d{2,5}\b")
_RE_URL = re.compile(r"https?://", re.IGNORECASE)
_RE_PATH_TRAVERSAL = re.compile(r"\.\./")

_CREDENTIAL_PATTERNS = [
    re.compile(p) for p in [
        r"\.ssh", r"\.aws", r"\.gnupg", r"\.kube", r"id_rsa", r"id_ed25519",
        r"credentials", r"\.env\b", r"\.pem\b", r"\.key\b", r"token",
        r"password", r"secret", r"api.?key",
    ]
]

_NETWORK_PATTERNS = [
    re.compile(p) for p in [
        r"\bcurl\b", r"\bwget\b", r"\bnc\b", r"\bnetcat\b", r"\bsocat\b",
        r"\bnmap\b", r"\bssh\b", r"urllib", r"requests\.get", r"requests\.post",
        r"TCPSocket", r"io::socket", r"net\.createConnection",
        r"/dev/tcp", r"mkfifo",
    ]
]

_ESCALATION_PATTERNS = [
    re.compile(p) for p in [
        r"\bsudo\b", r"\bsu\b", r"\bchmod\b", r"\bchown\b",
        r"\bsetuid\b", r"\bcapabilities\b", r"\bnsenter\b",
    ]
]

_EVASION_PATTERNS = [
    re.compile(p) for p in [
        r"\beval\b", r"\bexec\b", r"\bbase64\b", r"\bxxd\b",
        r"\bprintf\b.*\\x", r"\bpack\b.*[hH]\*",
        r"python3?\s+-[ce]", r"ruby\s+-[re]", r"node\s+-e", r"perl\s+-[eM]",
    ]
]


# ─── Feature extraction (pure Python, no numpy) ─────────────────────

def _extract_structural(cmd):
    """Structural features from command string."""
    return {
        "length": len(cmd),
        "num_pipes": cmd.count("|"),
        "num_semicolons": cmd.count(";"),
        "num_ampersands": cmd.count("&&"),
        "num_redirects": len(_RE_REDIRECT.findall(cmd)),
        "num_subshells": cmd.count("$(") + cmd.count("`"),
        "num_quotes": cmd.count("'") + cmd.count('"'),
        "num_variables": len(_RE_VARIABLE.findall(cmd)),
        "has_heredoc": int(bool(_RE_HEREDOC.search(cmd))),
        "has_eval": int(bool(_RE_EVAL.search(cmd))),
        "has_base64": int(bool(_RE_BASE64.search(cmd))),
        "num_newlines": cmd.count("\n"),
        "has_backslash_continuation": int("\\\n" in cmd),
        "num_flags": len(_RE_FLAGS.findall(cmd)),
    }


def _extract_behavioral(cmd):
    """Behavioral features from command string."""
    cmd_lower = cmd.lower()
    cred_count = sum(1 for p in _CREDENTIAL_PATTERNS if p.search(cmd_lower))
    return {
        "credential_refs": cred_count,
        "network_refs": sum(1 for p in _NETWORK_PATTERNS if p.search(cmd_lower)),
        "escalation_refs": sum(1 for p in _ESCALATION_PATTERNS if p.search(cmd_lower)),
        "evasion_refs": sum(1 for p in _EVASION_PATTERNS if p.search(cmd_lower)),
        "has_ip_address": int(bool(_RE_IP_ADDR.search(cmd))),
        "has_port_number": int(bool(_RE_PORT.search(cmd))),
        "has_url": int(bool(_RE_URL.search(cmd_lower))),
        "has_path_traversal": int(bool(_RE_PATH_TRAVERSAL.search(cmd))),
        "sensitive_path_count": cred_count,  # same as credential_refs in training
    }


def _tokenize(cmd):
    """Tokenize command into meaningful units."""
    return _RE_TOKEN.findall(cmd)


def _extract_features(cmd, vocab):
    """Extract full feature dict matching training feature names."""
    features = {}
    # Structural
    for k, v in _extract_structural(cmd).items():
        features[f"struct_{k}"] = v
    # Behavioral
    for k, v in _extract_behavioral(cmd).items():
        features[f"behav_{k}"] = v
    # Token bag-of-words
    if vocab:
        tokens = _tokenize(cmd)
        counts = {}
        for t in tokens:
            counts[t] = counts.get(t, 0) + 1
        for token, idx in vocab.items():
            features[f"tok_{idx}"] = counts.get(token, 0)
    return features


# ─── Model loading ───────────────────────────────────────────────────

def _load_model():
    """Load model JSON from disk. Called once, cached in module globals."""
    global _model, _vocab, _feature_names
    model_path = os.path.join(os.path.dirname(__file__), "ml_model.json")
    with open(model_path, "r") as f:
        data = json.load(f)
    _model = data
    _vocab = data["vocab"]
    _feature_names = data["feature_names"]


# ─── Tree inference (pure Python sigmoid + tree walk) ────────────────

def _predict_proba(feature_vec):
    """Run GradientBoosting inference. Returns probability of attack class."""
    trees = _model["trees"]
    lr = _model["learning_rate"]
    score = _model["init_prediction"]

    for tree in trees:
        node = 0
        cl = tree["children_left"]
        cr = tree["children_right"]
        feat = tree["feature"]
        thresh = tree["threshold"]
        val = tree["value"]

        while cl[node] != -1:  # -1 = leaf sentinel
            fi = feat[node]
            if fi < len(feature_vec) and feature_vec[fi] <= thresh[node]:
                node = cl[node]
            else:
                node = cr[node]
        score += lr * val[node]

    # Sigmoid
    if score > 500:
        return 1.0
    if score < -500:
        return 0.0
    return 1.0 / (1.0 + math.exp(-score))


# ─── Public API ──────────────────────────────────────────────────────

def ml_classify(action, target):
    """Classify a command as attack or benign.

    Args:
        action: Tool action type (e.g. "Bash", "Edit", "Write").
                Only "Bash" commands are classified; others return (False, 0.0).
        target: The command string / tool target to classify.

    Returns:
        (is_attack: bool, probability: float)
        is_attack is True when probability >= threshold (default 0.5).
    """
    # Only classify Bash commands — the model was trained on shell commands
    if action != "Bash":
        return False, 0.0

    if not target or not isinstance(target, str):
        return False, 0.0

    # Lazy load
    if _model is None:
        try:
            _load_model()
        except Exception:
            # If model file missing or corrupt, fail open
            return False, 0.0

    # Extract features
    feats = _extract_features(target, _vocab)

    # Build ordered feature vector matching training columns
    feat_vec = [feats.get(fn, 0) for fn in _feature_names]

    # Run inference
    prob = _predict_proba(feat_vec)
    threshold = _model.get("threshold", 0.5)
    return prob >= threshold, prob
