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 from request.state, with fallback to X-Tenant-ID 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.
"""

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          # <-- v4.3.1
from app.database.models_intents import DecisionAuditLogDB, TenantDB
from agentic_reliability_framework.core.models.event import ReliabilityEvent
from agentic_reliability_framework.core.governance.healing_intent import HealingIntent
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__)
router = APIRouter()


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


# --------------------------------------------------------------------------
# Helper: write audit log (idempotent)
# --------------------------------------------------------------------------
async def write_audit_log(
    db: Session,
    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.
    """
    # 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}")


# --------------------------------------------------------------------------
# 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
):
    """
    Evaluate an infrastructure intent with idempotency, tenant isolation,
    full governance loop analysis, and Bayesian skill posterior injection.
    """
    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 = request.headers.get("Authorization", "").replace("Bearer ", "")
    if not api_key:
        api_key = request.query_params.get("api_key", "unknown")

    # Get tenant_id from request.state or fallback to X-Tenant-ID header
    tenant_id = getattr(request.state, "tenant_id", None)
    if not tenant_id:
        tenant_id = request.headers.get("X-Tenant-ID")
    if not tenant_id:
        if span:
            span.set_status(Status(StatusCode.ERROR, "Missing tenant_id"))
            span.end()
        raise HTTPException(status_code=403, detail="Tenant not identified")

    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 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,
        )

        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,
            db=db,
            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),
):
    """Record an outcome for a previously evaluated intent."""
    try:
        risk_engine = request.app.state.risk_engine
        outcome_record = record_outcome(
            db=db,
            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 (unchanged)
# --------------------------------------------------------------------------
@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"),
):
    """
    Evaluate a healing decision, audit it, and optionally enforce via Rust ladder.
    """
    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 = request.headers.get("Authorization", "").replace("Bearer ", "")
    if not api_key:
        api_key = request.query_params.get("api_key", "unknown")

    tenant_id = getattr(request.state, "tenant_id", None)
    if not tenant_id:
        tenant_id = request.headers.get("X-Tenant-ID")
    if not tenant_id:
        if span:
            span.set_status(Status(StatusCode.ERROR, "Missing tenant_id"))
            span.end()
        raise HTTPException(status_code=403, detail="Tenant not identified")

    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,
        )

        # ---- 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,
            db=db,
            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)