app/services/risk_service.py

"""
Risk service – integrates ARF Bayesian risk engine, policy engine, and decision engine.
Deterministic, no random fallbacks, explicit error handling. Tenant‑aware.

Version: 2026-07-06 – added evaluate_intent_full with GovernanceLoop integration,
skill context injection, and full HealingIntent serialisation.
v4.3.1 – healing decision now optionally incorporates skill reliability
for Bayesian utility‑aware action selection.
"""

import json
import logging
import os
import time
from typing import Optional, List, Dict, Any

from agentic_reliability_framework.core.governance.risk_engine import RiskEngine
from agentic_reliability_framework.core.governance.intents import InfrastructureIntent
from agentic_reliability_framework.core.models.event import ReliabilityEvent, HealingAction
from agentic_reliability_framework.core.governance.policy_engine import PolicyEngine
from agentic_reliability_framework.core.decision.decision_engine import DecisionEngine
from agentic_reliability_framework.runtime.memory.rag_graph import RAGGraphMemory
from agentic_reliability_framework.core.research.eclipse_probe import compute_epistemic_risk

# ── Governance loop integration ──────────────────────────────
from agentic_reliability_framework.core.governance.governance_loop import GovernanceLoop
from agentic_reliability_framework.core.governance.cost_estimator import CostEstimator
from agentic_reliability_framework.core.governance.policies import PolicyEvaluator, allow_all
from agentic_reliability_framework.core.governance.stability_controller import LyapunovStabilityController
from agentic_reliability_framework.core.temporal_reliability import TemporalReliabilityMonitor
from agentic_reliability_framework.core.governance.healing_intent import HealingIntent

# ── optional tracing ─────────────────────────────────────────
try:
    from opentelemetry import trace
    _tracer = trace.get_tracer(__name__)
    OTEL_AVAILABLE = True
except ImportError:
    OTEL_AVAILABLE = False
    _tracer = None

# ── Prometheus metrics (always registered; no‑op if not scraped) ─
from prometheus_client import Counter, Histogram

_EVAL_COUNTER = Counter(
    "arf_evaluations_total",
    "Total evaluation calls (intent + healing), partitioned by engine and status.",
    ["engine", "status"],
)

_EVAL_DURATION = Histogram(
    "arf_evaluation_duration_seconds",
    "End‑to‑end latency of evaluation calls.",
    ["engine"],
    buckets=(0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0),
)

_RUST_AGREEMENT = Counter(
    "arf_rust_agreement_total",
    "Agreement between Rust enforcer and Python policy evaluation.",
    ["result"],  # "agreed" or "diverged"
)

# ── optional Rust enforcer (shadow mode) ──────────────────────
_RUST_ENFORCER_AVAILABLE = False
_rust_evaluator = None            # singleton per process
_rust_policy_json: Optional[str] = None

if os.getenv("ARF_USE_RUST_ENFORCER", "false").lower() == "true":
    try:
        import arf_enforcer
        _RUST_ENFORCER_AVAILABLE = True
    except ImportError:
        pass

# Default OSS policy tree – mirrors the hard‑coded rules in the Python PolicyEvaluator
_OSS_POLICY_TREE_JSON = json.dumps({
    "And": [
        {"Atomic": {"RegionAllowed": {"allowed_regions": ["eastus"]}}},
        {"Atomic": {"ResourceTypeRestricted": {
            "forbidden_types": ["DATABASE_DROP", "FULL_ROLLOUT", "SYSTEM_SHUTDOWN", "SECRET_ROTATION"]
        }}},
        {"Atomic": {"MaxPermissionLevel": {"max_level": "admin"}}}
    ]
})


def _ensure_rust_evaluator() -> bool:
    """Lazy initialise the Rust policy evaluator. Returns True on success."""
    global _rust_evaluator, _rust_policy_json
    if _rust_evaluator is not None:
        return True
    if not _RUST_ENFORCER_AVAILABLE:
        return False
    try:
        _rust_policy_json = _OSS_POLICY_TREE_JSON
        _rust_evaluator = arf_enforcer.PyPolicyEvaluator(_rust_policy_json)
        return True
    except Exception:
        _rust_evaluator = None
        return False


logger = logging.getLogger(__name__)


def evaluate_intent(
    engine: RiskEngine,
    intent: InfrastructureIntent,
    cost_estimate: Optional[float],
    policy_violations: List[str],
    tenant_id: Optional[str] = None,
) -> dict:
    """
    Evaluate an infrastructure intent using the Bayesian risk engine.

    The risk score is computed using a weighted fusion of conjugate online
    model, optional hyperpriors, and offline HMC. The tenant_id is passed
    to the risk engine to select the correct per‑tenant Beta store.

    Parameters
    ----------
    engine : RiskEngine
        Initialised ARF Bayesian risk engine (must be tenant‑aware).
    intent : InfrastructureIntent
        The infrastructure request to evaluate.
    cost_estimate : float or None
        Estimated monthly cost (used by cost‑threshold policies).
    policy_violations : list[str]
        Pre‑computed policy violation strings (from the Python evaluator).
    tenant_id : str, optional
        Tenant UUID. If provided, the risk engine will use tenant‑specific
        conjugate state. Required for multi‑tenant deployments.

    Returns
    -------
    dict
        Keys: risk_score, explanation, contributions.
    """
    t0 = time.monotonic()
    span = None
    if OTEL_AVAILABLE and _tracer:
        span = _tracer.start_span("risk_service.evaluate_intent")
        span.set_attribute("intent_type", type(intent).__name__)
        if tenant_id:
            span.set_attribute("tenant_id", tenant_id)

    # ── Shadow Rust enforcer (best‑effort, non‑blocking) ──────
    if _RUST_ENFORCER_AVAILABLE and _ensure_rust_evaluator():
        try:
            rust_intent = {
                "action": getattr(intent, "intent_type", "unknown"),
                "component": getattr(intent, "service_name", "unknown"),
                "region": getattr(intent, "region", None),
                "resource_type": getattr(intent, "resource_type", None),
                "permission_level": getattr(intent, "permission_level", None),
                "tenant_id": tenant_id,
                "extra": {}
            }
            rust_raw = _rust_evaluator.evaluate(
                json.dumps(rust_intent), cost_estimate
            )
            rust_violations = json.loads(rust_raw)

            agreed = set(rust_violations) == set(policy_violations)
            _RUST_AGREEMENT.labels(result="agreed" if agreed else "diverged").inc()
            if not agreed:
                msg = (
                    f"Rust enforcer divergence for tenant {tenant_id}: "
                    f"Rust={sorted(rust_violations)} Python={sorted(policy_violations)}"
                )
                logger.warning(msg)
                if span:
                    span.add_event("rust_enforcer_divergence", {
                        "rust_violations": rust_violations,
                        "python_violations": policy_violations
                    })
        except Exception as exc:
            logger.debug("Rust enforcer shadow evaluation failed: %s", exc)

    # ── Core risk evaluation ──────────────────────────────────
    try:
        if hasattr(engine, "set_tenant"):
            engine.set_tenant(tenant_id)
        elif tenant_id:
            logger.warning(
                "RiskEngine does not yet support tenant_id; evaluations will be shared across tenants."
            )

        score, explanation, contributions = engine.calculate_risk(
            intent=intent,
            cost_estimate=cost_estimate,
            policy_violations=policy_violations
        )
        engine_label = "python"
        status = "success"
    except Exception:
        _EVAL_COUNTER.labels(engine="python", status="error").inc()
        _EVAL_DURATION.labels(engine="python").observe(time.monotonic() - t0)
        raise

    _EVAL_COUNTER.labels(engine=engine_label, status=status).inc()
    _EVAL_DURATION.labels(engine=engine_label).observe(time.monotonic() - t0)

    if span:
        span.set_attribute("risk_score", score)
        if _RUST_ENFORCER_AVAILABLE:
            span.set_attribute("rust_enforcer_available", True)
        span.end()

    return {
        "risk_score": score,
        "explanation": explanation,
        "contributions": contributions
    }


def evaluate_intent_full(
    intent: InfrastructureIntent,
    *,
    risk_engine: RiskEngine,
    cost_estimator: Optional[CostEstimator] = None,
    policy_evaluator: Optional[PolicyEvaluator] = None,
    memory: Optional[RAGGraphMemory] = None,
    enable_epistemic: bool = False,
    hallucination_probe: Optional[Any] = None,
    predictive_engine: Optional[Any] = None,
    business_calculator: Optional[Any] = None,
    use_rust_enforcer: bool = False,
    stability_controller: Optional[LyapunovStabilityController] = None,
    temporal_monitor: Optional[TemporalReliabilityMonitor] = None,
    tenant_id: Optional[str] = None,
    skill_id: Optional[str] = None,
    skill_registry: Optional[Any] = None,
    context_extra: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Run the full governance loop and return a structured response containing
    the serialised HealingIntent with Bayesian skill posterior parameters.

    If stability_controller or temporal_monitor are None (the default),
    the governance loop will simply skip those checks. Pass stateful
    instances from the app state to accumulate cross‑request state.

    Parameters
    ----------
    intent : InfrastructureIntent
        The original infrastructure request.
    risk_engine : RiskEngine
        Bayesian risk engine (tenant‑aware).
    cost_estimator : CostEstimator, optional
        Monthly cost estimator; a default instance is created if None.
    policy_evaluator : PolicyEvaluator, optional
        Policy tree evaluator; defaults to `allow_all` if None.
    memory : RAGGraphMemory, optional
        Semantic memory for similar‑incident retrieval.
    enable_epistemic : bool
        Whether to run the ECLIPSE hallucination probe and CUDL attribution.
    hallucination_probe : HallucinationRisk, optional
        Pre‑configured probe instance.
    predictive_engine : SimplePredictiveEngine, optional
        Time‑series forecasting engine.
    business_calculator : BusinessImpactCalculator, optional
        Revenue impact estimator.
    use_rust_enforcer : bool
        Whether to run the Rust policy evaluator in shadow mode.
    stability_controller : LyapunovStabilityController, optional
        Passive stability monitor; if None, stability checks are skipped.
    temporal_monitor : TemporalReliabilityMonitor, optional
        Drift detector; if None, drift detection is skipped.
    tenant_id : str, optional
        Tenant UUID for multi‑tenant state.
    skill_id : str, optional
        Skill identifier; if provided, the skill's current posterior
        parameters are injected into the governance loop.
    skill_registry : SkillRegistry, optional
        Instance of the skill registry (required if skill_id is given).
    context_extra : dict, optional
        Additional key‑value pairs to merge into the loop context.

    Returns
    -------
    dict
        Keys:
        - risk_score : float
        - explanation : str
        - contributions : dict (empty; full trace is in healing_intent)
        - healing_intent : dict (serialised HealingIntent)
        - recommended_action : str
        - deterministic_id : str
    """
    t0 = time.monotonic()
    span = None
    if OTEL_AVAILABLE and _tracer:
        span = _tracer.start_span("risk_service.evaluate_intent_full")
        span.set_attribute("intent_type", type(intent).__name__)
        if tenant_id:
            span.set_attribute("tenant_id", tenant_id)

    # Default components if not provided
    if policy_evaluator is None:
        policy_evaluator = PolicyEvaluator(allow_all())
    if cost_estimator is None:
        cost_estimator = CostEstimator()
    # stability_controller and temporal_monitor are NOT defaulted here;
    # they remain None unless explicitly passed. The GovernanceLoop will skip
    # those checks gracefully.

    loop = GovernanceLoop(
        policy_evaluator=policy_evaluator,
        cost_estimator=cost_estimator,
        risk_engine=risk_engine,
        memory=memory,
        enable_epistemic=enable_epistemic,
        hallucination_probe=hallucination_probe,
        predictive_engine=predictive_engine,
        business_calculator=business_calculator,
        use_rust_enforcer=use_rust_enforcer,
        stability_controller=stability_controller,
        temporal_monitor=temporal_monitor,
    )

    # ── Build context with skill posterior parameters ─────────
    context: Dict[str, Any] = dict(context_extra) if context_extra else {}
    if skill_id and skill_registry is not None:
        try:
            # Fetch the latest version for the skill
            versions = skill_registry.list_skill_versions(skill_id)
            version = versions[-1] if versions else 1
            # Use public get_model() instead of direct _models access
            model = skill_registry.get_model(skill_id, version)
            if model is not None:
                alpha = model.alpha
                beta = model.beta
                reliability = model.mean()
            else:
                # Use default prior if no model exists yet
                alpha = skill_registry.default_prior_alpha
                beta = skill_registry.default_prior_beta
                reliability = alpha / (alpha + beta)
            context.update({
                "skill_id": skill_id,
                "skill_version": version,
                "skill_ate": skill_registry.get_ate(skill_id, version),
                "skill_reliability_score": reliability,
                "skill_alpha": alpha,
                "skill_beta": beta,
            })
        except Exception as e:
            logger.warning("Failed to inject skill context for '%s': %s", skill_id, e)

    # ── Execute governance loop ───────────────────────────────
    healing_intent: HealingIntent = loop.run(intent, context=context)
    healing_dict = healing_intent.to_dict(include_advisory_context=True)

    risk_score = healing_intent.risk_score or 0.0
    explanation = healing_intent.justification or ""

    # ── Metrics & span finalisation ───────────────────────────
    _EVAL_COUNTER.labels(engine="governance_loop", status="success").inc()
    _EVAL_DURATION.labels(engine="governance_loop").observe(time.monotonic() - t0)

    if span:
        span.set_attribute("risk_score", risk_score)
        span.set_attribute("recommended_action", healing_dict.get("recommended_action"))
        span.end()

    return {
        "risk_score": risk_score,
        "explanation": explanation,
        "contributions": {},  # full trace is in healing_intent
        "healing_intent": healing_dict,
        "recommended_action": healing_dict.get("recommended_action"),
        "deterministic_id": healing_intent.deterministic_id,
    }


def evaluate_healing_decision(
    event: ReliabilityEvent,
    policy_engine: PolicyEngine,
    decision_engine: Optional[DecisionEngine] = None,
    rag_graph: Optional[RAGGraphMemory] = None,
    model=None,
    tokenizer=None,
    tenant_id: Optional[str] = None,
    # ── v4.3.1: skill context ──────────────────────────────────
    skill_id: Optional[str] = None,
    skill_version: Optional[int] = None,
    skill_registry: Optional[Any] = None,
) -> Dict[str, Any]:
    """
    Evaluate healing actions for a given reliability event using decision‑theoretic selection.
    Includes epistemic risk signals from the eclipse probe and, optionally, skill reliability
    information to bias the utility towards actions from trusted skills.

    The utility of each candidate action a is extended with two additional terms:

        U(a) = U_base(a) + w_skill · μ_skill − w_σ · σ_skill

    where μ_skill = α/(α+β) is the posterior mean reliability of the skill that
    authored the action, and σ_skill = sqrt(αβ / ((α+β)²(α+β+1))) is its
    posterior standard deviation.  These terms are computed from the conjugate
    Beta posterior tracked by the SkillRegistry.  When no skill context is
    provided, the utility falls back to the original formulation.

    Parameters
    ----------
    event : ReliabilityEvent
        The incident event containing latency, error rate, etc.
    policy_engine : PolicyEngine
        The ARF healing policy engine with configured policies.
    decision_engine : DecisionEngine, optional
        If omitted, a default instance is created. If provided, it is used as‑is
        (its internal skill registry is not modified).
    rag_graph : RAGGraphMemory, optional
        Semantic memory for similar incident retrieval.
    model, tokenizer : optional
        HuggingFace model and tokenizer for epistemic risk computation.
    tenant_id : str, optional
        Tenant UUID for logging and metrics.
    skill_id : str, optional
        Skill identifier to incorporate into utility.
    skill_version : int, optional
        Version of the skill.
    skill_registry : SkillRegistry, optional
        Registry to fetch the skill's posterior parameters.

    Returns
    -------
    dict
        Keys: risk_score, selected_action, expected_utility, alternatives,
        explanation, epistemic_signals, plus skill_id/skill_version if present.
    """
    t0 = time.monotonic()
    span = None
    if OTEL_AVAILABLE and _tracer:
        span = _tracer.start_span("risk_service.evaluate_healing")
        span.set_attribute("component", event.component)
        if tenant_id:
            span.set_attribute("tenant_id", tenant_id)

    # If decision_engine not provided, try to get from policy_engine
    if decision_engine is None and hasattr(policy_engine, 'decision_engine'):
        decision_engine = policy_engine.decision_engine

    # If still None, create a minimal one (global stats only), passing skill registry if available
    if decision_engine is None:
        logger.debug("No DecisionEngine provided; creating default instance")
        decision_engine = DecisionEngine(rag_graph=rag_graph, skill_registry=skill_registry)

    # Get raw candidate actions (by temporarily disabling decision engine)
    orig_use = policy_engine.use_decision_engine
    try:
        policy_engine.use_decision_engine = False
        raw_actions = policy_engine.evaluate_policies(event)
    finally:
        policy_engine.use_decision_engine = orig_use

    # If no actions, return NO_ACTION
    if not raw_actions or raw_actions == [HealingAction.NO_ACTION]:
        if span:
            span.set_attribute("selected_action", HealingAction.NO_ACTION.value)
            span.end()
        _EVAL_COUNTER.labels(engine="python", status="success").inc()
        _EVAL_DURATION.labels(engine="python").observe(time.monotonic() - t0)
        return {
            "risk_score": 0.0,
            "selected_action": HealingAction.NO_ACTION.value,
            "expected_utility": 0.0,
            "alternatives": [],
            "explanation": "No candidate actions triggered.",
            "epistemic_signals": None,
        }

    # Build reasoning text from policies that triggered the actions
    reasoning_parts = []
    for policy in policy_engine.policies:
        if any(a in policy.actions for a in raw_actions):
            conditions_str = ", ".join(
                f"{c.metric} {c.operator} {c.threshold}" for c in policy.conditions
            )
            reasoning_parts.append(
                f"Policy {policy.name} triggered by {conditions_str} → actions {[a.value for a in policy.actions]}"
            )
    reasoning_text = " ".join(reasoning_parts)

    # Build evidence text from the event
    evidence_text = (
        f"Component: {event.component}, "
        f"latency_p99: {event.latency_p99}, "
        f"error_rate: {event.error_rate}, "
        f"cpu_util: {event.cpu_util}, "
        f"memory_util: {event.memory_util}"
    )

    # Compute epistemic signals (if model/tokenizer provided)
    epistemic_signals = None
    if model is not None and tokenizer is not None:
        try:
            epistemic_signals = compute_epistemic_risk(
                reasoning_text, evidence_text, model, tokenizer
            )
        except Exception as e:
            logger.error(f"Failed to compute epistemic risk: {e}")
            epistemic_signals = {
                "entropy": 0.0,
                "contradiction": 0.0,
                "evidence_lift": 0.0,
                "hallucination_risk": 0.0,
            }
    else:
        logger.debug("Epistemic model/tokenizer not provided; using zero signals")
        epistemic_signals = {
            "entropy": 0.0,
            "contradiction": 0.0,
            "evidence_lift": 0.0,
            "hallucination_risk": 0.0,
        }

    # ── Decision with skill context ──────────────────────────
    decision = decision_engine.select_optimal_action(
        raw_actions,
        event,
        component=event.component,
        epistemic_signals=epistemic_signals,
        skill_id=skill_id,
        skill_version=skill_version,
    )

    # Extract risk of the selected action
    risk_score = None
    for alt in decision.alternatives:
        if alt.action == decision.best_action:
            risk_score = alt.risk
            break
    if risk_score is None:
        # Compute risk separately
        risk_score = decision_engine.compute_risk(
            decision.best_action, event, event.component)

    # Format alternatives (top 3 only)
    alt_list = []
    for alt in decision.alternatives[:3]:
        alt_list.append({
            "action": alt.action.value,
            "expected_utility": alt.utility,
            "risk": alt.risk,
        })

    # ── Metrics & span finalisation ───────────────────────────
    _EVAL_COUNTER.labels(engine="python", status="success").inc()
    _EVAL_DURATION.labels(engine="python").observe(time.monotonic() - t0)

    if span:
        span.set_attribute("risk_score", risk_score)
        span.set_attribute("selected_action", decision.best_action.value)
        span.set_attribute("expected_utility", decision.expected_utility)
        span.end()

    result = {
        "risk_score": risk_score,
        "selected_action": decision.best_action.value,
        "expected_utility": decision.expected_utility,
        "alternatives": alt_list,
        "explanation": decision.explanation,
        "raw_decision": decision.raw_data,
        "epistemic_signals": epistemic_signals,
    }
    if skill_id:
        result["skill_id"] = skill_id
        result["skill_version"] = skill_version
    return result


def get_system_risk() -> float:
    """
    Return an aggregated risk score across all monitored components.
    This endpoint is deprecated. Use component‑level risk evaluation instead.
    """
    raise NotImplementedError(
        "get_system_risk is deprecated. Use component‑level risk evaluation instead."
    )