"""Tier-1 HTTP model at inference: hashed byte n-grams -> int8 embedding bag -> small MLP, numpy only.

File (.npz): emb int8 [B, E] + scale per column, W1 [E+S, H] / b1, W2 [H, C] / b2 (float32), info json
(buckets, orders, side-feature names, classes, temperature, thresholds). The same hashing (dibaguard.hashing)
runs at training time inside torch's EmbeddingBag, so bucket ids agree bit for bit. Side features (encoding depth,
charset profile, lengths) ride along the bag. One request costs one pass of ngram_buckets + two small matmuls.
"""
from __future__ import annotations

import json

import numpy as np

from dibaguard.hashing import bag_mean, ngram_buckets
from dibaguard.http.canon import charset_profile

SIDE = ["deepest", "len", "n_params", "persian", "non_ascii", "special", "digits"]


def side_features(texts: list[str], deepest) -> np.ndarray:
    out = np.zeros((len(texts), len(SIDE)), np.float32)
    for i, (t, d) in enumerate(zip(texts, deepest)):
        pr = charset_profile(t)
        out[i] = (float(d), min(len(t), 8192) / 1024.0, t.count(" | "), pr.get("persian", 0.0), pr.get("non_ascii", 0.0),
                  pr.get("special", 0.0), pr.get("digits", 0.0))
    return out


class HttpModel:
    def __init__(self, emb_q: np.ndarray, scale: np.ndarray, W1, b1, W2, b2, info: dict):
        self.emb_q, self.scale = np.asarray(emb_q, np.int8), np.asarray(scale, np.float32)
        self.W1, self.b1, self.W2, self.b2 = (np.asarray(a, np.float32) for a in (W1, b1, W2, b2))
        self.info = info
        self.buckets, self.orders = int(info["buckets"]), tuple(info["orders"])
        self.max_bytes = int(info.get("max_bytes", 1024))

    def embed(self, texts: list[str]) -> np.ndarray:
        idx, doc = ngram_buckets(texts, self.buckets, self.orders, self.max_bytes)
        return bag_mean(self.emb_q, idx, doc, len(texts), self.scale)

    def logits(self, texts: list[str], deepest=None) -> np.ndarray:
        v = self.embed(texts)
        if self.W1.shape[0] > v.shape[1]:
            side = side_features(texts, deepest if deepest is not None else [0] * len(texts))
            v = np.hstack([v, side])
        h = np.maximum(v @ self.W1 + self.b1, 0.0)
        return h @ self.W2 + self.b2

    def proba_attack(self, texts: list[str], deepest=None) -> np.ndarray:
        z = self.logits(texts, deepest)[:, 0] / float(self.info.get("temperature", 1.0))
        return 1.0 / (1.0 + np.exp(-z))

    def save(self, path: str) -> None:
        np.savez_compressed(path, emb=self.emb_q, scale=self.scale, W1=self.W1, b1=self.b1, W2=self.W2, b2=self.b2,
                            info=np.frombuffer(json.dumps(self.info).encode(), np.uint8))

    @classmethod
    def load(cls, path: str) -> "HttpModel":
        z = np.load(path)
        return cls(z["emb"], z["scale"], z["W1"], z["b1"], z["W2"], z["b2"], json.loads(bytes(z["info"]).decode()))

    @staticmethod
    def quantise(emb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """float [B, E] -> (int8, per-column scale) so that int8 * scale ~ emb."""
        scale = np.maximum(np.abs(emb).max(0), 1e-8) / 127.0
        return np.clip(np.round(emb / scale), -127, 127).astype(np.int8), scale.astype(np.float32)


def novelty_scores(texts, seen: np.ndarray, buckets: int, orders, max_bytes: int = 1024, leave_one_out: bool = False) -> np.ndarray:
    idx, doc = ngram_buckets(texts, buckets, orders, max_bytes)
    out = np.zeros(len(texts))
    if idx.size == 0:
        return out
    cnt = seen[idx].astype(np.float64)
    if leave_one_out:                                 # subtract this request's own contribution to the vocabulary
        key = doc.astype(np.int64) * buckets + idx
        _, inv, own = np.unique(key, return_inverse=True, return_counts=True)
        cnt = np.maximum(cnt - own[inv], 0.0)
    w = 1.0 / (1.0 + cnt)                             # unseen -> 1, seen once -> 0.5, common -> ~0
    tot = np.bincount(doc, minlength=len(texts)).astype(np.float64)
    return np.bincount(doc, weights=w, minlength=len(texts)) / np.maximum(tot, 1.0)


class HttpModel2(HttpModel):
    """Site-conditioned tier-1 (http_v2): the same bag + MLP, plus a *site profile* taken from a sample of the site's
    own benign requests - the mean bag vector (a request's cosine distance to it is an input feature) and the mean
    hidden vector, which modulates the hidden layer (FiLM). A site without a profile uses the one shipped from
    training. Weights: W1 [E+S+1, H], Wc [H, 2H] / bc, W2 [H, C]."""

    def __init__(self, emb_q, scale, W1, b1, Wc, bc, W2, b2, info: dict, site: dict | None = None):
        super().__init__(emb_q, scale, W1, b1, W2, b2, info)
        self.Wc, self.bc = np.asarray(Wc, np.float32), np.asarray(bc, np.float32)
        self.site = site or {}

    def _hidden(self, texts, deepest, s_mean: np.ndarray) -> np.ndarray:
        v = self.embed(texts)
        side = side_features(texts, deepest if deepest is not None else [0] * len(texts))
        num = (v * s_mean[None, :]).sum(1)
        den = np.linalg.norm(v, axis=1) * np.linalg.norm(s_mean) + 1e-6
        dist = (1.0 - num / den)[:, None].astype(np.float32)
        return np.maximum(np.hstack([v, side, dist]) @ self.W1 + self.b1, 0.0)

    def site_profile(self, texts: list[str], deepest=None) -> dict:
        """Benign requests of a site (a few hundred suffice) -> {'s_mean': [E], 'ctx': [H], 'seen': bucket counts [B]}.
        'seen' is the site's own n-gram vocabulary: what this application's normal requests are made of."""
        v = self.embed(texts)
        s_mean = v.mean(0).astype(np.float32)
        idx, _ = ngram_buckets(texts, self.buckets, self.orders, self.max_bytes)
        seen = np.bincount(idx, minlength=self.buckets).astype(np.int32) if idx.size else np.zeros(self.buckets, np.int32)
        return {"s_mean": s_mean, "ctx": self._hidden(texts, deepest, s_mean).mean(0).astype(np.float32), "seen": seen, "n_ref": len(texts)}

    def novelty(self, texts: list[str], site: dict | None = None, leave_one_out: bool = False) -> np.ndarray:
        """Site-only tier: how much of a request is made of n-grams this site's normal traffic never contains, weighted
        by rarity (an n-gram seen in one reference request counts almost as new) -> [n] in [0, 1]. No labels.
        leave_one_out: the texts are the reference itself - each request's own n-grams are taken out of the counts,
        so reference and live requests are scored the same way (exchangeable, as the threshold needs)."""
        site = site or self.site
        seen = site.get("seen")
        if seen is None:
            return np.zeros(len(texts))
        return novelty_scores(texts, seen, self.buckets, self.orders, self.max_bytes, leave_one_out)

    def logits(self, texts: list[str], deepest=None, site: dict | None = None) -> np.ndarray:
        site = site or self.site
        h = self._hidden(texts, deepest, site["s_mean"])
        g = np.tanh(site["ctx"] @ self.Wc + self.bc)
        H = h.shape[1]
        h = h * (1.0 + g[:H]) + g[H:]
        return h @ self.W2 + self.b2

    def proba_attack(self, texts: list[str], deepest=None, site: dict | None = None) -> np.ndarray:
        z = self.logits(texts, deepest, site)[:, 0] / float(self.info.get("temperature", 1.0))
        return 1.0 / (1.0 + np.exp(-z))

    def save(self, path: str) -> None:
        np.savez_compressed(path, emb=self.emb_q, scale=self.scale, W1=self.W1, b1=self.b1, Wc=self.Wc, bc=self.bc, W2=self.W2,
                            b2=self.b2, site_s_mean=self.site["s_mean"], site_ctx=self.site["ctx"],
                            site_seen=self.site.get("seen", np.zeros(self.buckets, np.int32)),
                            info=np.frombuffer(json.dumps(self.info).encode(), np.uint8))

    @classmethod
    def load(cls, path: str) -> "HttpModel2":
        z = np.load(path)
        return cls(z["emb"], z["scale"], z["W1"], z["b1"], z["Wc"], z["bc"], z["W2"], z["b2"],
                   json.loads(bytes(z["info"]).decode()), {"s_mean": z["site_s_mean"], "ctx": z["site_ctx"],
                                                           "seen": z["site_seen"] if "site_seen" in z.files else None})