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)
- Auditable decision logging (DecisionAuditLogDB)
- Pricing telemetry (optional, to arf‑pricing‑calculator)
- OpenTelemetry tracing
- Optional Rust execution ladder for mechanical enforcement
"""

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
from typing import Optional, Dict, Any

from app.models.infrastructure_intents import InfrastructureIntentRequest
from app.services.intent_adapter import to_oss_intent
from app.services.risk_service import evaluate_intent, 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
from app.database.models_intents import DecisionAuditLogDB, TenantDB  # <-- NEW
from agentic_reliability_framework.core.models.event import ReliabilityEvent
from agentic_reliability_framework.core.governance.healing_intent import HealingIntent

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


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 from HealingIntent (or result dict)
    risk_score = healing_intent.get("risk_score", 0.5)
    action = healing_intent.get("recommended_action", "deny")  # approve/deny/escalate
    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", [])

    # Memory and causal fields (usually in metadata)
    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")  # could be derived from risk_factors
    causal_effect = metadata.get("causal_effect")

    # Build audit entry
    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,  # same as action for now
        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}")


# --------------------------------------------------------------------------
# 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"),
):
    """
    Evaluate an infrastructure intent with idempotency, tenant isolation, and audit logging.
    """
    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 (set by enforce_quota)
    tenant_id = getattr(request.state, "tenant_id", None)
    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

        # TODO: Modify risk_service.evaluate_intent to accept tenant_id
        # and pass it down to RiskEngine (which will select the correct BetaStore)
        result = evaluate_intent(
            engine=risk_engine,
            intent=oss_intent,
            cost_estimate=intent_req.estimated_cost,
            policy_violations=intent_req.policy_violations,
            # tenant_id=tenant_id   # after modification
        )

        if span:
            span.set_attribute("risk_score", result["risk_score"])
            span.set_attribute("deterministic_id", str(uuid.uuid4()))

        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,
            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) ----
        # Extract the HealingIntent dictionary from result (if not present, construct minimal)
        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 (idempotent, pricing)
# --------------------------------------------------------------------------
@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"),
):
    """
    Record an outcome for a previously evaluated intent.
    Also updates the pricing calculator's calibration buffer if available.
    """
    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,
        )

        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 (with optional Rust enforcement)
# --------------------------------------------------------------------------
@router.post("/healing/evaluate")
async def evaluate_healing_decision_endpoint(
    request: Request,
    decision_req: HealingDecisionRequest,
    background_tasks: BackgroundTasks,
    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:
        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:
                # Convert response_data to a HealingIntent dict (or use the actual HealingIntent object)
                # For simplicity, assume response_data contains the same fields as HealingIntent.to_enterprise_request()
                intent_dict = response_data.get("healing_intent", response_data)
                ladder = ExecutionLadder()
                rust_result = ladder.evaluate(intent_dict)
                if not rust_result.get("allowed", False):
                    # Override decision
                    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 ----
        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)