app/api/routes_governance.py

from fastapi import APIRouter, Depends, HTTPException, Request, BackgroundTasks
from fastapi.encoders import jsonable_encoder
from sqlalchemy.orm import Session
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 pydantic import BaseModel
import uuid
import logging
import time

from agentic_reliability_framework.core.models.event import ReliabilityEvent

# ===== USAGE TRACKER IMPORTS =====
from app.core.usage_tracker import enforce_quota, UsageRecord, tracker

logger = logging.getLogger(__name__)
router = APIRouter()


class OutcomeRequest(BaseModel):
    deterministic_id: str
    success: bool
    recorded_by: str
    notes: str = ""


class HealingDecisionRequest(BaseModel):
    event: ReliabilityEvent


@router.post("/intents/evaluate")
async def evaluate_intent_endpoint(
    request: Request,
    intent_req: InfrastructureIntentRequest,
    background_tasks: BackgroundTasks,
    db: Session = Depends(get_db),
    quota: dict = Depends(enforce_quota)
):
    start_time = time.time()
    api_key = quota["api_key"]
    tier = quota["tier"]
    response_data = None
    error_msg = None

    try:
        oss_intent = to_oss_intent(intent_req)
        risk_engine = request.app.state.risk_engine
        result = evaluate_intent(
            engine=risk_engine,
            intent=oss_intent,
            cost_estimate=intent_req.estimated_cost,
            policy_violations=intent_req.policy_violations
        )

        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

        if tracker:
            record = UsageRecord(
                api_key=api_key,
                tier=tier,
                timestamp=time.time(),
                endpoint="/api/v1/intents/evaluate",
                request_body=intent_req.model_dump(),
                response=response_data,
                processing_ms=(time.time() - start_time) * 1000,
            )
            await tracker.increment_usage_async(record, background_tasks)

        return response_data

    except HTTPException:
        raise
    except Exception as e:
        error_msg = str(e)
        logger.exception("Error in evaluate_intent_endpoint")
        if tracker:
            record = UsageRecord(
                api_key=api_key,
                tier=tier,
                timestamp=time.time(),
                endpoint="/api/v1/intents/evaluate",
                request_body=intent_req.model_dump(),
                error=error_msg,
                processing_ms=(time.time() - start_time) * 1000,
            )
            await tracker.increment_usage_async(record, background_tasks)
        raise HTTPException(status_code=500, detail=error_msg)


@router.post("/intents/outcome")
async def record_outcome_endpoint(
    request: Request,
    outcome: OutcomeRequest,
    db: Session = Depends(get_db)
):
    # No usage tracking for outcomes (doesn't count against quota)
    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
        )
        return {"message": "Outcome recorded", "outcome_id": outcome_record.id}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/healing/evaluate")
async def evaluate_healing_decision_endpoint(
    request: Request,
    decision_req: HealingDecisionRequest,
    background_tasks: BackgroundTasks,
    quota: dict = Depends(enforce_quota)
):
    start_time = time.time()
    api_key = quota["api_key"]
    tier = quota["tier"]
    response_data = None
    error_msg = None

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

        if tracker:
            record = UsageRecord(
                api_key=api_key,
                tier=tier,
                timestamp=time.time(),
                endpoint="/api/v1/healing/evaluate",
                request_body=decision_req.model_dump(),
                response=response_data,
                processing_ms=(time.time() - start_time) * 1000,
            )
            await tracker.increment_usage_async(record, background_tasks)

        return response_data

    except HTTPException:
        raise
    except Exception as e:
        error_msg = str(e)
        logger.exception("Error in evaluate_healing_decision_endpoint")
        if tracker:
            record = UsageRecord(
                api_key=api_key,
                tier=tier,
                timestamp=time.time(),
                endpoint="/api/v1/healing/evaluate",
                request_body=decision_req.model_dump(),
                error=error_msg,
                processing_ms=(time.time() - start_time) * 1000,
            )
            await tracker.increment_usage_async(record, background_tasks)
        raise HTTPException(status_code=500, detail=error_msg)