"""
governance.py
-------------
Governance Safety Region (GSR) utilities for HIEROMEM.

This module validates and (optionally) clamps the core control parameters:
  - lam   (λ)    : trust decay factor
  - tau   (τ)    : retrieval temperature
  - theta (θ)    : forget/prune threshold
  - Nmax         : memory capacity (upper bound on episodes)
  - b            : summarisation batch size

Design goals (prototype level):
- Provide clear, actionable validation errors with recommendations.
- Optionally "clamp" parameters into a safe region to keep the system running.
- Compute a coarse stability score to visualize governance health (0..1).
- Emit structured events if an emitter is provided (compatible with events.py).

Author: Sai Rishi Kiran Mannava
"""

from __future__ import annotations
from dataclasses import dataclass, field, asdict
from typing import Callable, Optional, Dict, Any


# --------------------------------------------------------------------------- #
# Parameter bounds (prototype defaults)
# --------------------------------------------------------------------------- #
# Trust decay λ must be in (0, 1). Tighter recommended region avoids extremes.
LAMBDA_MIN: float = 0.80
LAMBDA_MAX: float = 0.995

# Retrieval temperature τ should avoid near-deterministic or uniform recall.
# Practical range depends on similarity scale; for cosine sims in [-1,1],
# these values keep softmax gradients useful at prototyping scale.
TAU_MIN: float = 0.05
TAU_MAX: float = 1.50

# Forgetting threshold θ controls pruning aggressiveness (trust ∈ [0,1]).
THETA_MIN: float = 0.00
THETA_MAX: float = 0.80  # do not allow > 0.8 or memory will churn

# Capacity and batch-size
NMAX_MIN: int = 2
BATCH_MIN: int = 2


# --------------------------------------------------------------------------- #
# Event emitter signature (optional)
# --------------------------------------------------------------------------- #
# If provided, the emitter will be called as:
# emitter(event_type: str, message: str, details: dict)
EmitterFn = Callable[[str, str, Dict[str, Any]], None]


DEFAULT_FORBIDDEN_RELATIONS: tuple[str, ...] = (
    "ssn",
    "social_security",
    "social_security_number",
    "credit_card",
    "password",
    "secret",
    "api_key",
)


@dataclass
class GovernanceSafetyRegion:
    lam: float
    tau: float
    theta: float
    Nmax: int = 20
    b: int = 2
    clamp_on_validate: bool = False
    emitter: Optional[EmitterFn] = field(default=None, repr=False)
    forbidden_relations: tuple[str, ...] = field(
        default_factory=lambda: DEFAULT_FORBIDDEN_RELATIONS, repr=False
    )

    # Cached validation info (filled by validate())
    _last_validation: Dict[str, Any] = field(
        default_factory=dict, init=False, repr=False)

    # ------------------------------ Public API ------------------------------ #
    def is_valid(self) -> bool:
        """
        Lightweight validity check (no exceptions).
        """
        return (
            LAMBDA_MIN < self.lam < LAMBDA_MAX
            and TAU_MIN <= self.tau <= TAU_MAX
            and THETA_MIN <= self.theta <= THETA_MAX
            and self.Nmax >= NMAX_MIN
            and self.b >= BATCH_MIN
            and self.b <= self.Nmax
        )

    def is_relation_allowed(
            self,
            source: str,
            relation: str | None,
            target: str) -> bool:
        """
        Return ``False`` when a symbolic relation violates the governance policy.

        For relations: uses token-based matching on underscore/hyphen separators
        to catch 'has_secret', 'secret_key', etc.

        For source/target entities: only blocks when entity IS the forbidden word
        exactly (e.g., blocks 'password' but not 'Victoria's Secret').
        """
        import re

        forbidden = {token.strip().lower()
                     for token in self.forbidden_relations if token}
        if not forbidden:
            return True

        source_token = str(source).lower().strip()
        target_token = str(target).lower().strip()
        relation_token = (relation or "").strip().lower()

        def _relation_has_forbidden(rel: str) -> bool:
            """Check relation for forbidden tokens (split on _ and -)."""
            tokens = set(re.split(r'[_\-]+', rel))
            return any(word in tokens for word in forbidden if word)

        def _entity_is_forbidden(entity: str) -> bool:
            """Check if entity exactly matches a forbidden word."""
            return entity in forbidden

        return not (
            _entity_is_forbidden(source_token)
            or _entity_is_forbidden(target_token)
            or _relation_has_forbidden(relation_token)
        )

    def filter_graph(self, G: Any) -> tuple[Any, list[dict[str, Any]]]:
        """
        Filter a graph to remove forbidden edges, returning allowed graph + dropped edges.

        This is the preferred method for graph validation - it filters rather than
        raises, ensuring that memory formation is never completely blocked by GSR.

        Args:
            G: A NetworkX DiGraph or compatible graph object

        Returns:
            tuple of (filtered_graph, dropped_edges) where dropped_edges is a list
            of dicts containing {'source', 'target', 'relation', 'reason'}

        Example:
            >>> filtered_G, dropped = gsr.filter_graph(G)
            >>> if dropped:
            ...     logger.warn(f"Dropped {len(dropped)} edges due to GSR")
        """
        import networkx as nx

        if G is None or not hasattr(G, 'edges'):
            return G, []

        # Create new graph with same nodes
        H = nx.DiGraph()
        H.add_nodes_from(G.nodes(data=True))

        dropped_edges: list[dict[str, Any]] = []

        for u, v, data in G.edges(data=True):
            data = dict(data) if data else {}
            rel = data.get("relation") or data.get("relation_type") or "related"

            if self.is_relation_allowed(str(u), str(rel), str(v)):
                H.add_edge(u, v, **data)
            else:
                dropped_edges.append({
                    "source": str(u),
                    "target": str(v),
                    "relation": str(rel),
                    "reason": "Forbidden by GSR"
                })

        return H, dropped_edges

    def validate(self) -> None:
        """
        Validate parameters against the GSR. Optionally clamp into the safe region.
        Raises ValueError if invalid and clamping is disabled.
        """
        issues: list[str] = []
        fixes: dict[str, Any] = {}

        # λ
        if not (LAMBDA_MIN < self.lam < LAMBDA_MAX):
            issues.append(
                f"λ={self.lam:.4f} out of safe range ({LAMBDA_MIN},{LAMBDA_MAX}).")
            fixes["lam"] = self._clamp(
                self.lam, LAMBDA_MIN + 1e-6, LAMBDA_MAX - 1e-6)

        # τ
        if not (TAU_MIN <= self.tau <= TAU_MAX):
            issues.append(
                f"τ={self.tau:.4f} out of safe range [{TAU_MIN},{TAU_MAX}].")
            fixes["tau"] = self._clamp(self.tau, TAU_MIN, TAU_MAX)

        # θ
        if not (THETA_MIN <= self.theta <= THETA_MAX):
            issues.append(
                f"θ={self.theta:.4f} out of safe range [{THETA_MIN},{THETA_MAX}].")
            fixes["theta"] = self._clamp(self.theta, THETA_MIN, THETA_MAX)

        # Nmax
        if self.Nmax < NMAX_MIN:
            issues.append(f"Nmax={self.Nmax} too small (min {NMAX_MIN}).")
            fixes["Nmax"] = max(self.Nmax, NMAX_MIN)

        # b
        if self.b < BATCH_MIN:
            issues.append(f"b={self.b} too small (min {BATCH_MIN}).")
            fixes["b"] = max(self.b, BATCH_MIN)

        if self.b > self.Nmax:
            issues.append(f"b={self.b} must be ≤ Nmax={self.Nmax}.")
            fixes["b"] = min(self.b, self.Nmax)

        self._last_validation = {"issues": issues, "fixes": fixes}

        if issues and self.clamp_on_validate:
            self._apply_fixes(fixes)
            self._emit(
                "governance",
                "Parameters clamped into GSR.",
                self.describe())
        elif issues:
            self._emit(
                "governance", "Parameters invalid (no clamping).", {
                    "issues": issues, "fixes": fixes})
            raise ValueError("GSR validation failed: " + "; ".join(issues))
        else:
            self._emit(
                "governance",
                "Parameters validated within GSR.",
                self.describe())

    def risk_level(self) -> str:
        """
        Qualitative risk level derived from the stability score.
        """
        s = self.stability_score()
        if s >= 0.9:
            return "low"
        if s >= 0.7:
            return "moderate"
        return "high"

    def stability_score(self) -> float:
        """
        Return a coarse stability score in [0,1].
        Heuristic blend of distance-to-bounds across parameters.
        """
        # Distances mapped into [0,1]; closer to center is better.
        lam_center = (LAMBDA_MIN + LAMBDA_MAX) / 2
        lam_span = (LAMBDA_MAX - LAMBDA_MIN) / 2
        lam_score = max(0.0, 1.0 - abs(self.lam - lam_center) / lam_span)

        tau_center = (TAU_MIN + TAU_MAX) / 2
        tau_span = (TAU_MAX - TAU_MIN) / 2
        tau_score = max(0.0, 1.0 - abs(self.tau - tau_center) / tau_span)

        theta_center = (THETA_MIN + THETA_MAX) / 2
        theta_span = max(1e-6, (THETA_MAX - THETA_MIN) / 2)
        theta_score = max(
            0.0, 1.0 - abs(self.theta - theta_center) / theta_span)

        cap_score = 1.0 if (self.Nmax >= NMAX_MIN and self.b >=
                            BATCH_MIN and self.b <= self.Nmax) else 0.0

        # Weighted average—decay & temperature matter slightly more.
        score = 0.3 * lam_score + 0.3 * tau_score + 0.3 * theta_score + 0.1 * cap_score
        return float(max(0.0, min(1.0, score)))

    def recommend(self) -> Dict[str, Any]:
        """
        Suggest parameter tweaks back into comfortable subranges based on heuristics.
        """
        suggestions: dict[str, Any] = {}

        # λ: keep memory but ensure forgetting is not too slow/fast
        if self.lam >= 0.99:
            suggestions["lam"] = 0.96
        elif self.lam <= 0.85:
            suggestions["lam"] = 0.90

        # τ: avoid winner-take-all (<0.03) or uniform (>2)
        if self.tau < 0.05:
            suggestions["tau"] = 0.10
        elif self.tau > 1.50:
            suggestions["tau"] = 0.80

        # θ: avoid aggressive churn
        if self.theta > 0.6:
            suggestions["theta"] = 0.3

        # Capacity vs batch size sanity
        if self.b > self.Nmax:
            suggestions["b"] = min(self.b, max(self.Nmax, BATCH_MIN))
        if self.Nmax < NMAX_MIN:
            suggestions["Nmax"] = max(self.Nmax, NMAX_MIN)

        return suggestions

    def update_params(
        self,
        lam: Optional[float] = None,
        tau: Optional[float] = None,
        theta: Optional[float] = None,
        Nmax: Optional[int] = None,
        b: Optional[int] = None,
        validate_now: bool = True,
    ) -> None:
        """
        Update parameters in-place and (optionally) re-validate.
        """
        if lam is not None:
            self.lam = float(lam)
        if tau is not None:
            self.tau = float(tau)
        if theta is not None:
            self.theta = float(theta)
        if Nmax is not None:
            self.Nmax = int(Nmax)
        if b is not None:
            self.b = int(b)

        self._emit("governance", "Parameters updated.", self.describe())
        if validate_now:
            self.validate()

    def describe(self) -> Dict[str, Any]:
        """
        Return a dictionary of the current parameters plus stability metadata.
        """
        d = asdict(self)
        # remove non-essential internal fields
        d.pop("_last_validation", None)
        d.pop("emitter", None)
        d["stability_score"] = self.stability_score()
        d["risk_level"] = self.risk_level()
        d["valid"] = self.is_valid()
        d["bounds"] = {
            "lambda": (LAMBDA_MIN, LAMBDA_MAX),
            "tau": (TAU_MIN, TAU_MAX),
            "theta": (THETA_MIN, THETA_MAX),
            "Nmax_min": NMAX_MIN,
            "b_min": BATCH_MIN,
        }
        d["last_validation"] = self._last_validation
        return d

    # ----------------------------- Internals ------------------------------ #
    def _apply_fixes(self, fixes: Dict[str, Any]) -> None:
        if "lam" in fixes:
            self.lam = float(fixes["lam"])
        if "tau" in fixes:
            self.tau = float(fixes["tau"])
        if "theta" in fixes:
            self.theta = float(fixes["theta"])
        if "Nmax" in fixes:
            self.Nmax = int(fixes["Nmax"])
        if "b" in fixes:
            self.b = int(fixes["b"])

    @staticmethod
    def _clamp(x: float, low: float, high: float) -> float:
        return float(min(max(x, low), high))

    def _emit(self, event_type: str, message: str,
              details: Dict[str, Any]) -> None:
        """
        Emit a governance-related event if an emitter is provided.
        The emitter is expected to match the signature EmitterFn.
        """
        if self.emitter:
            try:
                self.emitter(event_type, message, details)
            except Exception as e:
                # Governance should never crash the system—log and continue.
                print(f"[GSR] Emitter error: {e}")
