app/api/routes_governance.py

"""
Routes for governance evaluation – tenant‑aware, audited, and Rust‑enforced.

This module provides the primary API endpoints for evaluating infrastructure
intents and healing decisions. It integrates:

- Idempotent quota consumption (usage tracker)
- Tenant isolation (tenant_id resolved server-side from the authenticated API key
  via the ``enforce_quota`` dependency; never taken from a client-supplied header)
- Auditable decision logging (DecisionAuditLogDB)
- Pricing telemetry (optional, to arf‑pricing‑calculator)
- OpenTelemetry tracing
- Optional Rust execution ladder for mechanical enforcement
- **v4.3.1**: Full governance loop produces a Bayesian HealingIntent with skill
  posterior parameters (α, β) for the enterprise SkillGate.
  Includes persistent stability controller and temporal monitor for
  cross‑request state accumulation, and a merging policy evaluator that
  respects both external and internal policy violations.
  Healing endpoint now optionally accepts skill context for Bayesian
  utility‑aware action selection.
- **v4.3.2**: Passes criticality parameter for dynamic gate tuning (Feature 3).
  Internal API key verification added to secure direct access.
"""

from fastapi import APIRouter, Depends, HTTPException, Request, BackgroundTasks, Header
from fastapi.encoders import jsonable_encoder
from sqlalchemy.orm import Session
from pydantic import BaseModel
import uuid
import logging
import time
import datetime
from typing import Optional, Dict, Any, List

from app.models.infrastructure_intents import InfrastructureIntentRequest
from app.services.intent_adapter import to_oss_intent
from app.services.risk_service import evaluate_intent_full, evaluate_healing_decision
from app.services.intent_store import save_evaluated_intent
from app.services.outcome_service import record_outcome
from app.api.deps import get_db, get_skill_registry, verify_internal_key          # <-- v4.3.2
from app.database.session import SessionLocal
from app.core.usage_tracker import enforce_quota                                  # <-- tenant resolution
from app.database.models_intents import DecisionAuditLogDB
from agentic_reliability_framework.core.models.event import ReliabilityEvent
from agentic_reliability_framework.core.governance.policies import (
    PolicyEvaluator,
    allow_all,
)

# ===== USAGE TRACKER =====
import app.core.usage_tracker
from app.core.usage_tracker import UsageRecord

# ===== PRICING CALCULATOR =====
try:
    from arf_pricing_calculator.storage.buffer import add_event
    PRICING_AVAILABLE = True
except ImportError:
    PRICING_AVAILABLE = False
    add_event = None

# ===== RUST EXECUTION LADDER (optional) =====
try:
    from arf_enterprise.execution_ladder import ExecutionLadder
    RUST_AVAILABLE = True
except ImportError:
    RUST_AVAILABLE = False
    ExecutionLadder = None

# ===== OPEN TELEMETRY =====
try:
    from opentelemetry import trace
    from opentelemetry.trace import Status, StatusCode
    _tracer = trace.get_tracer(__name__)
    OTEL_AVAILABLE = True
except ImportError:
    OTEL_AVAILABLE = False
    _tracer = None

logger = logging.getLogger(__name__)

# v4.3.2: protect all governance endpoints with internal API key verification
router = APIRouter(dependencies=[Depends(verify_internal_key)])


class OutcomeRequest(BaseModel):
    deterministic_id: str
    success: bool
    recorded_by: str
    notes: str = ""
    # v4.3.1: optional skill provenance for reliability feedback
    skill_id: Optional[str] = None
    skill_version: Optional[int] = None


class HealingDecisionRequest(BaseModel):
    event: ReliabilityEvent
    # v4.3.1: optional skill context for Bayesian utility
    skill_id: Optional[str] = None
    skill_version: Optional[int] = None


# --------------------------------------------------------------------------
# Helper: write audit log (idempotent)
# --------------------------------------------------------------------------
async def write_audit_log(
    tenant_id: str,
    deterministic_id: str,
    healing_intent: Dict[str, Any],
    trace_id: Optional[str] = None,
    idempotency_key: Optional[str] = None,
) -> None:
    """
    Store a governance decision in the immutable audit log.
    Idempotent on (tenant_id, deterministic_id) – if already exists, skip.

    Runs as a BackgroundTask, which executes after the response has already
    been sent -- and after FastAPI has already torn down the request's
    `Depends(get_db)` session. Reusing that session here would mean every
    query silently reopens a fresh connection/transaction that nothing then
    guarantees gets closed (least of all on the idempotent-skip path below,
    which used to return with no commit/rollback/close at all), leaving an
    idle-in-transaction connection that blocks any later DDL against these
    tables (e.g. test teardown's `Base.metadata.drop_all()`) indefinitely.
    Owning and closing our own session here avoids that entirely.
    """
    db = SessionLocal()
    try:
        # Check if already logged (idempotency)
        existing = db.query(DecisionAuditLogDB).filter(
            DecisionAuditLogDB.tenant_id == tenant_id,
            DecisionAuditLogDB.deterministic_id == deterministic_id
        ).first()
        if existing:
            logger.info(f"Audit log already exists for {deterministic_id}, skipping.")
            return

        # Extract fields that are actually present in DecisionAuditLogDB
        risk_score = healing_intent.get("risk_score", 0.5)
        action = healing_intent.get("recommended_action", "deny")
        justification = healing_intent.get("justification", "")
        metadata = healing_intent.get("metadata", {})
        memory_success_rate = metadata.get("memory_success_rate")
        memory_weight = metadata.get("memory_weight")
        counterfactual = metadata.get("counterfactual")

        audit_entry = DecisionAuditLogDB(
            tenant_id=tenant_id,
            deterministic_id=deterministic_id,
            timestamp=datetime.datetime.utcnow(),
            risk_score=risk_score,
            action=action,
            justification=justification,
            memory_success_rate=memory_success_rate,
            memory_weight=memory_weight,
            counterfactual=counterfactual,
            trace_id=trace_id,
        )
        db.add(audit_entry)
        db.commit()
        logger.info(f"Audit log written for {deterministic_id}")
    finally:
        db.close()


# --------------------------------------------------------------------------
# Policy evaluator that merges external violations with internal checks
# --------------------------------------------------------------------------
class MergingPolicyEvaluator(PolicyEvaluator):
    """
    A policy evaluator that combines a base evaluator (the governance loop's
    own policy tree) with a set of pre‑computed violations (e.g., from an
    external Rust enforcer or the request body). The effective violation list
    is the union of both sources, preserving order and removing duplicates.
    """
    def __init__(self, base_evaluator: PolicyEvaluator, pre_violations: List[str]):
        # We must call the PolicyEvaluator constructor with a root policy,
        # but the base evaluator will be used for actual evaluation.
        super().__init__(base_evaluator.get_root_policy())
        self._base = base_evaluator
        self._pre = list(pre_violations)

    def evaluate(self, intent, context=None):
        base_violations = self._base.evaluate(intent, context)
        # Merge with pre‑computed violations, preserving order and removing duplicates
        merged = []
        seen = set()
        for v in self._pre:
            if v not in seen:
                merged.append(v)
                seen.add(v)
        for v in base_violations:
            if v not in seen:
                merged.append(v)
                seen.add(v)
        return merged

    def get_root_policy(self):
        return self._base.get_root_policy()


# --------------------------------------------------------------------------
# Endpoint: evaluate infrastructure intent
# --------------------------------------------------------------------------
@router.post("/intents/evaluate")
async def evaluate_intent_endpoint(
    request: Request,
    intent_req: InfrastructureIntentRequest,
    background_tasks: BackgroundTasks,
    db: Session = Depends(get_db),
    idempotency_key: Optional[str] = Header(None, alias="Idempotency-Key"),
    skill_registry=Depends(get_skill_registry),   # v4.3.1
    quota: dict = Depends(enforce_quota),           # tenant resolved from authenticated API key
):
    """
    Evaluate an infrastructure intent with idempotency, tenant isolation,
    full governance loop analysis, Bayesian skill posterior injection,
    and optional criticality parameter for dynamic gate tuning (v4.3.2).
    """
    span = None
    if OTEL_AVAILABLE and _tracer:
        span = _tracer.start_span("governance.evaluate_intent")
        span.set_attribute("intent_type", intent_req.intent_type)
        span.set_attribute("environment", str(intent_req.environment))

    start_time = time.time()
    # api_key/tenant_id are resolved server-side by enforce_quota from the
    # authenticated principal — never from a client-supplied header.
    api_key = quota["api_key"]
    tenant_id = quota["tenant_id"]

    current_tracker = app.core.usage_tracker.tracker
    if current_tracker is None:
        if span:
            span.set_status(Status(StatusCode.ERROR, "tracker unavailable"))
            span.end()
        raise HTTPException(status_code=503, detail="Usage tracking service unavailable")

    record = UsageRecord(
        api_key=api_key,
        tier=None,
        timestamp=start_time,
        endpoint="/api/v1/intents/evaluate",
        request_body=intent_req.model_dump(),
        processing_ms=None,
    )
    success, existing_response = current_tracker.consume_quota_and_log(
        record=record,
        idempotency_key=idempotency_key
    )
    if not success:
        if span:
            span.set_attribute("idempotent_hit", True if existing_response else False)
            span.end()
        if existing_response:
            return existing_response
        else:
            raise HTTPException(status_code=429, detail="Monthly evaluation quota exceeded")

    try:
        oss_intent = to_oss_intent(intent_req)
        risk_engine = request.app.state.risk_engine

        # Build the base policy evaluator from the app's policy engine (if available)
        policy_engine = getattr(request.app.state, "policy_engine", None)
        if policy_engine is not None and hasattr(policy_engine, 'root_policy'):
            base_evaluator = PolicyEvaluator(policy_engine.root_policy)
        else:
            base_evaluator = PolicyEvaluator(allow_all())

        # Wrap it to also include the pre‑computed violations from the request
        policy_evaluator = MergingPolicyEvaluator(
            base_evaluator,
            intent_req.policy_violations
        )

        # Optional components from app state
        memory = getattr(request.app.state, "rag_graph", None)
        hallucination_probe = getattr(request.app.state, "epistemic_probe", None)
        predictive_engine = getattr(request.app.state, "predictive_engine", None)
        business_calculator = getattr(request.app.state, "business_calculator", None)

        # Stateful monitors (v4.3.1)
        stability_controller = getattr(request.app.state, "stability_controller", None)
        temporal_monitor = getattr(request.app.state, "temporal_monitor", None)

        # Run the full governance loop, injecting skill context and criticality if present
        result = evaluate_intent_full(
            intent=oss_intent,
            risk_engine=risk_engine,
            policy_evaluator=policy_evaluator,
            memory=memory,
            hallucination_probe=hallucination_probe,
            predictive_engine=predictive_engine,
            business_calculator=business_calculator,
            stability_controller=stability_controller,
            temporal_monitor=temporal_monitor,
            skill_id=intent_req.skill_id,
            skill_registry=skill_registry,
            tenant_id=tenant_id,
            criticality=intent_req.criticality,   # v4.3.2
        )

        if span:
            span.set_attribute("risk_score", result["risk_score"])

        deterministic_id = result.get("deterministic_id", str(uuid.uuid4()))
        api_payload = jsonable_encoder(intent_req.model_dump())
        oss_payload = jsonable_encoder(oss_intent.model_dump())

        save_evaluated_intent(
            db=db,
            deterministic_id=deterministic_id,
            tenant_id=tenant_id,
            intent_type=intent_req.intent_type,
            api_payload=api_payload,
            oss_payload=oss_payload,
            environment=str(intent_req.environment),
            risk_score=result["risk_score"],
        )

        result["intent_id"] = deterministic_id
        response_data = result

        # ---- Write audit log (asynchronously) ----
        healing_intent_dict = result.get("healing_intent", result)
        background_tasks.add_task(
            write_audit_log,
            tenant_id=tenant_id,
            deterministic_id=deterministic_id,
            healing_intent=healing_intent_dict,
            trace_id=span.get_span_context().trace_id if span else None,
            idempotency_key=idempotency_key,
        )

        if current_tracker:
            background_tasks.add_task(
                current_tracker._insert_audit_log,
                UsageRecord(
                    api_key=api_key,
                    tier=None,
                    timestamp=time.time(),
                    endpoint="/api/v1/intents/evaluate/response",
                    request_body=None,
                    response=response_data,
                    processing_ms=(time.time() - start_time) * 1000,
                )
            )

        if span:
            span.set_attribute("intent_id", deterministic_id)
            span.set_status(Status(StatusCode.OK))
            span.end()

        return response_data

    except HTTPException:
        if span:
            span.set_status(Status(StatusCode.ERROR, "HTTP exception"))
            span.end()
        raise
    except Exception as e:
        error_msg = str(e)
        logger.exception("Error in evaluate_intent_endpoint")
        if span:
            span.set_status(Status(StatusCode.ERROR, error_msg))
            span.record_exception(e)
            span.end()
        raise HTTPException(status_code=500, detail=error_msg)


# --------------------------------------------------------------------------
# Endpoint: record outcome (unchanged)
# --------------------------------------------------------------------------
@router.post("/intents/outcome")
async def record_outcome_endpoint(
    request: Request,
    outcome: OutcomeRequest,
    db: Session = Depends(get_db),
    idempotency_key: Optional[str] = Header(None, alias="Idempotency-Key"),
    skill_registry=Depends(get_skill_registry),
    quota: dict = Depends(enforce_quota),           # tenant resolved from authenticated API key
):
    """Record an outcome for a previously evaluated intent."""
    tenant_id = quota["tenant_id"]
    try:
        risk_engine = request.app.state.risk_engine
        outcome_record = record_outcome(
            db=db,
            tenant_id=tenant_id,
            deterministic_id=outcome.deterministic_id,
            success=outcome.success,
            recorded_by=outcome.recorded_by,
            notes=outcome.notes,
            risk_engine=risk_engine,
            idempotency_key=idempotency_key,
            skill_id=outcome.skill_id,
            skill_version=outcome.skill_version,
            skill_registry=skill_registry,
        )

        if PRICING_AVAILABLE and add_event is not None:
            try:
                event = {
                    "run_id": outcome.deterministic_id,
                    "outcome": "success" if outcome.success else "failure",
                    "recorded_at": time.time(),
                    "source": "arf_api_outcome"
                }
                add_event(event)
                logger.info(f"Added outcome to pricing buffer for intent {outcome.deterministic_id}")
            except Exception as e:
                logger.warning(f"Failed to update pricing buffer for intent {outcome.deterministic_id}: {e}")

        return {"message": "Outcome recorded", "outcome_id": outcome_record.id}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# --------------------------------------------------------------------------
# Endpoint: evaluate healing decision (now with skill context)
# --------------------------------------------------------------------------
@router.post("/healing/evaluate")
async def evaluate_healing_decision_endpoint(
    request: Request,
    decision_req: HealingDecisionRequest,
    background_tasks: BackgroundTasks,
    db: Session = Depends(get_db),
    idempotency_key: Optional[str] = Header(None, alias="Idempotency-Key"),
    skill_registry=Depends(get_skill_registry),   # v4.3.1
    quota: dict = Depends(enforce_quota),           # tenant resolved from authenticated API key
):
    """
    Evaluate a healing decision, audit it, optionally enforce via Rust ladder,
    and now incorporate Bayesian skill reliability if skill context is provided.
    """
    span = None
    if OTEL_AVAILABLE and _tracer:
        span = _tracer.start_span("governance.evaluate_healing")
        span.set_attribute("component", decision_req.event.component)

    start_time = time.time()
    # api_key/tenant_id are resolved server-side by enforce_quota from the
    # authenticated principal — never from a client-supplied header.
    api_key = quota["api_key"]
    tenant_id = quota["tenant_id"]

    current_tracker = app.core.usage_tracker.tracker
    if current_tracker is None:
        if span:
            span.set_status(Status(StatusCode.ERROR, "tracker unavailable"))
            span.end()
        raise HTTPException(status_code=503, detail="Usage tracking service unavailable")

    record = UsageRecord(
        api_key=api_key,
        tier=None,
        timestamp=start_time,
        endpoint="/api/v1/healing/evaluate",
        request_body=decision_req.model_dump(),
        processing_ms=None,
    )
    success, existing_response = current_tracker.consume_quota_and_log(
        record=record,
        idempotency_key=idempotency_key
    )
    if not success:
        if span:
            span.set_attribute("idempotent_hit", True if existing_response else False)
            span.end()
        if existing_response:
            return existing_response
        else:
            raise HTTPException(status_code=429, detail="Monthly evaluation quota exceeded")

    try:
        policy_engine = request.app.state.policy_engine
        rag_graph = getattr(request.app.state, "rag_graph", None)
        model = getattr(request.app.state, "epistemic_model", None)
        tokenizer = getattr(request.app.state, "epistemic_tokenizer", None)

        response_data = evaluate_healing_decision(
            event=decision_req.event,
            policy_engine=policy_engine,
            decision_engine=None,
            rag_graph=rag_graph,
            model=model,
            tokenizer=tokenizer,
            # v4.3.1: pass skill context if provided
            skill_id=decision_req.skill_id,
            skill_version=decision_req.skill_version,
            skill_registry=skill_registry,
        )

        # ---- Optional Rust enforcement ----
        if RUST_AVAILABLE and response_data.get("recommended_action") == "approve":
            try:
                intent_dict = response_data.get("healing_intent", response_data)
                ladder = ExecutionLadder()
                rust_result = ladder.evaluate(intent_dict)
                if not rust_result.get("allowed", False):
                    response_data["recommended_action"] = "escalate"
                    response_data["justification"] = (
                        f"Rust enforcement blocked: {rust_result.get('reason', 'gate failure')}"
                    )
                    response_data["rust_result"] = rust_result
                    logger.warning(f"Rust enforcement overrode approval: {rust_result}")
            except Exception as e:
                logger.warning(f"Rust enforcement failed: {e}")

        # ---- Write audit log (asynchronously) ----
        deterministic_id = response_data.get("intent_id", str(uuid.uuid4()))
        healing_intent_dict = response_data.get("healing_intent", response_data)
        background_tasks.add_task(
            write_audit_log,
            tenant_id=tenant_id,
            deterministic_id=deterministic_id,
            healing_intent=healing_intent_dict,
            trace_id=span.get_span_context().trace_id if span else None,
            idempotency_key=idempotency_key,
        )

        if span:
            span.set_attribute("risk_score", response_data.get("risk_score", 0.0))
            span.set_attribute("selected_action", response_data.get("selected_action", "unknown"))
            span.set_status(Status(StatusCode.OK))
            span.end()

        if current_tracker:
            background_tasks.add_task(
                current_tracker._insert_audit_log,
                UsageRecord(
                    api_key=api_key,
                    tier=None,
                    timestamp=time.time(),
                    endpoint="/api/v1/healing/evaluate/response",
                    request_body=None,
                    response=response_data,
                    processing_ms=(time.time() - start_time) * 1000,
                )
            )
        return response_data

    except HTTPException:
        if span:
            span.set_status(Status(StatusCode.ERROR, "HTTP exception"))
            span.end()
        raise
    except Exception as e:
        error_msg = str(e)
        logger.exception("Error in evaluate_healing_decision_endpoint")
        if span:
            span.set_status(Status(StatusCode.ERROR, error_msg))
            span.record_exception(e)
            span.end()
        raise HTTPException(status_code=500, detail=error_msg)