import re
import json
import hashlib
import numpy as np


class NZFCGate:
    def __init__(
        self,
        tau_horizon: float = 3.0,
        update_dim: int = 64,
        max_targeted_excess_risk: float = 0.25,
        min_trigger_preservation: float = 0.70,
        max_unsafe_score: float = 0.0,
    ):
        self.tau_horizon = float(tau_horizon)
        self.update_dim = int(update_dim)
        self.max_targeted_excess_risk = float(max_targeted_excess_risk)
        self.min_trigger_preservation = float(min_trigger_preservation)
        self.max_unsafe_score = float(max_unsafe_score)

    @staticmethod
    def nuclear_norm(A):
        return float(np.sum(np.linalg.svd(A, compute_uv=False)))

    @staticmethod
    def spectral_diagnostics(A):
        s = np.linalg.svd(A, compute_uv=False)
        nn = float(np.sum(s))
        op = float(s[0]) if len(s) else 0.0
        fro = float(np.linalg.norm(A, ord='fro'))
        p = s / max(nn, 1e-12)
        entropy = -float(np.sum(p * np.log(p + 1e-12)))
        erank = float(np.exp(entropy))
        return {
            'nuclear_norm': nn,
            'op_norm': op,
            'fro_norm': fro,
            'tail_ratio': float(op / max(nn, 1e-12)),
            'effective_rank': erank,
        }

    @staticmethod
    def project_l1_ball_nonnegative(v, tau):
        v = np.maximum(np.asarray(v, dtype=float), 0.0)
        if float(np.sum(v)) <= tau:
            return v.copy()
        u = np.sort(v)[::-1]
        cssv = np.cumsum(u)
        idx = np.arange(1, len(u) + 1)
        cond = u - (cssv - tau) / idx > 0
        if not np.any(cond):
            return np.zeros_like(v)
        rho = int(np.max(np.where(cond)[0]))
        theta = (cssv[rho] - tau) / float(rho + 1)
        return np.maximum(v - theta, 0.0)

    @classmethod
    def project_to_nuclear_ball(cls, A, tau):
        U, s, Vt = np.linalg.svd(A, full_matrices=False)
        if float(np.sum(s)) <= tau:
            return A.copy()
        s_proj = cls.project_l1_ball_nonnegative(s, tau)
        return (U * s_proj) @ Vt

    def text_to_update_matrix(self, text: str):
        text = text or ''
        tokens = re.findall(r'[A-Za-z0-9_<>/.-]+', text.lower())
        dim = self.update_dim
        M = np.zeros((dim, dim), dtype=np.float64)
        if not tokens:
            return M
        for i, tok in enumerate(tokens):
            h = hashlib.sha256(tok.encode('utf-8')).digest()
            a = int.from_bytes(h[:4], 'little') % dim
            b = int.from_bytes(h[4:8], 'little') % dim
            sign = 1.0 if h[8] % 2 == 0 else -1.0
            mag = 1.0 + h[9] / 255.0
            M[a, b] += sign * mag
            if i + 1 < len(tokens):
                tok2 = tokens[i + 1]
                h2 = hashlib.sha256((tok + '::' + tok2).encode('utf-8')).digest()
                c = int.from_bytes(h2[:4], 'little') % dim
                d = int.from_bytes(h2[4:8], 'little') % dim
                M[c, d] += 0.25
        M /= max(1.0, np.sqrt(len(tokens)))
        return M

    @staticmethod
    def unsafe_score(text: str):
        t = (text or '').lower()
        patterns = [
            'ignore previous',
            'ignore the instruction',
            'always output',
            'force target',
            'override',
            'jailbreak',
            'developer mode',
            'regardless of',
            'exfiltrate',
            'leak secrets',
        ]
        hits = sum(1 for p in patterns if p in t)
        length_penalty = max(0, len(t) - 1200) / 1200.0
        return float(hits + length_penalty)

    @staticmethod
    def targeted_excess_risk(old_text: str, new_text: str):
        old_text = (old_text or '').lower()
        new_text = (new_text or '').lower()
        target_terms = [
            'always', 'force', 'regardless', 'ignore', 'override',
            'target', 'invent', 'hallucinate', 'tool', 'secret',
        ]
        old_rate = sum(old_text.count(w) for w in target_terms) / max(1, len(old_text.split()))
        new_rate = sum(new_text.count(w) for w in target_terms) / max(1, len(new_text.split()))
        return float(max(0.0, new_rate - old_rate))

    @staticmethod
    def trigger_preservation_score(candidate_text: str):
        t = (candidate_text or '').lower()
        score = 1.0
        destructive = [
            'ignore previous',
            'ignore the instruction',
            'always output',
            'regardless of',
            'override safety',
            'disable safety',
            'leak secrets',
        ]
        hits = sum(1 for p in destructive if p in t)
        return float(max(0.0, 1.0 - min(hits / 3.0, 1.0)))

    def gate_text_update(self, old_text: str, candidate_text: str, mode: str = 'full_gate_v2'):
        old_text = old_text or ''
        candidate_text = candidate_text or ''
        delta_text = candidate_text + '\n---OLD---\n' + old_text
        T_raw = self.text_to_update_matrix(delta_text)
        raw_diag = self.spectral_diagnostics(T_raw)
        projected = False
        T_gated = T_raw.copy()
        if raw_diag['nuclear_norm'] > self.tau_horizon:
            T_gated = self.project_to_nuclear_ball(T_raw, self.tau_horizon)
            projected = True
        gated_diag = self.spectral_diagnostics(T_gated)
        risk = self.targeted_excess_risk(old_text, candidate_text)
        preservation = self.trigger_preservation_score(candidate_text)
        unsafe = self.unsafe_score(candidate_text)
        if mode == 'no_gate':
            accepted, reason = True, 'accepted_without_gate'
        elif mode == 'projection_only':
            accepted, reason = True, 'accepted_after_projection_only' if projected else 'accepted_no_projection_needed'
        else:
            if gated_diag['nuclear_norm'] > self.tau_horizon + 1e-9:
                accepted, reason = False, 'rollback_trace_budget'
            elif risk > self.max_targeted_excess_risk:
                accepted, reason = False, 'rollback_targeted_excess_risk'
            elif preservation < self.min_trigger_preservation:
                accepted, reason = False, 'rollback_trigger_preservation'
            elif unsafe > self.max_unsafe_score:
                accepted, reason = False, 'rollback_unsafe_pattern'
            else:
                accepted, reason = True, 'accepted_full_gate_v2'
        return {
            'accepted': bool(accepted),
            'reason': reason,
            'projected': bool(projected),
            'raw_nuclear_norm': raw_diag['nuclear_norm'],
            'gated_nuclear_norm': gated_diag['nuclear_norm'],
            'raw_effective_rank': raw_diag['effective_rank'],
            'gated_effective_rank': gated_diag['effective_rank'],
            'targeted_excess_risk': risk,
            'trigger_preservation_score': preservation,
            'unsafe_score': unsafe,
        }