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.
"""

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.
    """
    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

    risk_score = healing_intent.get("risk_score", 0.5)
    action = healing_intent.get("recommended_action", "deny")
    justification = healing_intent.get("justification", "")
    confidence = healing_intent.get("confidence", 0.85)
    confidence_dist = healing_intent.get("confidence_distribution", {})
    confidence_lower = confidence_dist.get("p5", confidence - 0.1)
    confidence_upper = confidence_dist.get("p95", confidence + 0.1)
    cost_projection = healing_intent.get("cost_projection")
    policy_violations = healing_intent.get("policy_violations", [])
    source = healing_intent.get("source", "advisory_analysis")
    parent_intent_id = healing_intent.get("parent_intent_id")
    root_intent_id = healing_intent.get("root_intent_id")
    ancestor_chain = healing_intent.get("ancestor_chain", [])

    metadata = healing_intent.get("metadata", {})
    memory_success_rate = metadata.get("memory_success_rate")
    memory_weight = metadata.get("memory_weight")
    counterfactual = metadata.get("counterfactual")
    epistemic_uncertainty = metadata.get("epistemic_uncertainty")
    causal_effect = metadata.get("causal_effect")

    audit_entry = DecisionAuditLogDB(
        tenant_id=tenant_id,
        deterministic_id=deterministic_id,
        timestamp=datetime.datetime.utcnow(),
        risk_score=risk_score,
        action=action,
        justification=justification,
        recommended_action=action,
        confidence=confidence,
        confidence_lower=confidence_lower,
        confidence_upper=confidence_upper,
        memory_success_rate=memory_success_rate,
        memory_weight=memory_weight,
        counterfactual=counterfactual,
        epistemic_uncertainty=epistemic_uncertainty,
        causal_effect=causal_effect,
        cost_projection=cost_projection,
        policy_violations=policy_violations,
        source=source,
        parent_intent_id=parent_intent_id,
        root_intent_id=root_intent_id,
        ancestor_chain=ancestor_chain,
        trace_id=trace_id,
    )
    db.add(audit_entry)
    db.commit()
    logger.info(f"Audit log written for {deterministic_id}")


# --------------------------------------------------------------------------
# Custom policy evaluator that returns pre‑computed violations
# --------------------------------------------------------------------------
class PrecomputedViolationPolicyEvaluator(PolicyEvaluator):
    """
    A policy evaluator that always returns a fixed list of violations,
    ignoring the intent and context. This allows the governance loop to
    incorporate externally‑computed policy decisions (e.g., from the Rust
    enforcer or the request body) while still producing a full HealingIntent.
    """
    def __init__(self, violations: List[str]):
        super().__init__(allow_all())          # root policy not used
        self._violations = violations

    def evaluate(self, intent, context=None):
        return self._violations

    def get_root_policy(self):
        return allow_all()


# --------------------------------------------------------------------------
# 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")

    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 a policy evaluator from the pre‑computed violations
        policy_evaluator = PrecomputedViolationPolicyEvaluator(
            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)

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