smart-org/apps/api/app/infrastructure/repositories/audit.py

"""Repositories for Audit & Observability and the Usage Ledger."""

from __future__ import annotations

from datetime import datetime

from sqlalchemy import case, func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.application.ports.audit import UsageAggregate, UsageBreakdown, UsageRecord
from app.domain.ai_control.models import RateUnit
from app.domain.audit.models import AuditEvent, ErrorEvent, QualityMetric, Span, Trace
from app.domain.base import new_id
from app.infrastructure.db.models import (
    AuditEventRow,
    ErrorEventRow,
    QualityMetricRow,
    SpanRow,
    TraceRow,
    UsageEventRow,
)
from app.infrastructure.repositories.credit import CreditLedgerRepository


def _audit_to_domain(row: AuditEventRow) -> AuditEvent:
    return AuditEvent(
        id=row.id,
        organization_id=row.organization_id,
        actor_user_id=row.actor_user_id,
        action=row.action,
        target_type=row.target_type,
        target_id=row.target_id,
        metadata_json=row.metadata_json or {},
        ip_address=row.ip_address,
        user_agent=row.user_agent,
        created_at=row.created_at,
    )


def _error_to_domain(row: ErrorEventRow) -> ErrorEvent:
    return ErrorEvent(
        id=row.id,
        organization_id=row.organization_id,
        run_id=row.run_id,
        kind=row.kind,
        message=row.message,
        stack=row.stack,
        metadata_json=row.metadata_json or {},
        source=row.source,
        severity=row.severity,
        error_type=row.error_type,
        fingerprint=row.fingerprint,
        status_code=row.status_code,
        method=row.method,
        path=row.path,
        user_id=row.user_id,
        request_id=row.request_id,
        ip=row.ip,
        user_agent=row.user_agent,
        resolved_at=row.resolved_at,
        resolved_note=row.resolved_note,
        created_at=row.created_at,
    )


class AuditRepository:
    """Concrete ``AuditPort``. Append-only."""

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def record(self, *, event: AuditEvent) -> AuditEvent:
        row = AuditEventRow(
            id=event.id,
            organization_id=event.organization_id,
            actor_user_id=event.actor_user_id,
            action=event.action,
            target_type=event.target_type,
            target_id=event.target_id,
            metadata_json=event.metadata_json,
            ip_address=event.ip_address,
            user_agent=event.user_agent,
            created_at=event.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _audit_to_domain(row)

    async def list_for_organization(
        self, *, organization_id: str, limit: int = 100, offset: int = 0
    ) -> list[AuditEvent]:
        stmt = (
            select(AuditEventRow)
            .where(AuditEventRow.organization_id == organization_id)
            .order_by(AuditEventRow.created_at.desc())
            .limit(min(limit, 500))
            .offset(offset)
        )
        return [_audit_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]

    async def list_platform(self, *, limit: int = 100, offset: int = 0) -> list[AuditEvent]:
        stmt = (
            select(AuditEventRow)
            .order_by(AuditEventRow.created_at.desc())
            .limit(min(limit, 500))
            .offset(offset)
        )
        return [_audit_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]


class UsageLedgerRepository:
    """Concrete ``UsageLedgerPort``. The ledger is the billing source of truth.

    Prepaid: every usage row also debits its organization's wallet, once, in the
    same transaction (the debit is keyed on the usage row's id). The ledger stays
    the source of truth; the wallet is the balance it leaves behind.
    """

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def record(self, *, usage: UsageRecord) -> float | None:
        """Append one usage row, debit the wallet for it, return the new balance.

        Both writes share the caller's transaction, so a rolled-back request leaves
        neither behind. Returns ``None`` when the row cost nothing and nothing was
        debited.
        """
        total = usage.total_price or (usage.quantity * usage.unit_price)
        row = UsageEventRow(
            id=new_id(),
            organization_id=usage.organization_id,
            project_id=usage.project_id,
            user_id=usage.user_id,
            employee_id=usage.employee_id,
            run_id=usage.run_id,
            unit=usage.unit.value,
            quantity=usage.quantity,
            unit_price=usage.unit_price,
            internal_cost_estimate=usage.internal_cost_estimate,
            total_price=total,
            metadata_json=usage.metadata_json,
        )
        self.session.add(row)
        await self.session.flush()
        debit = await CreditLedgerRepository(self.session).debit_for_usage(
            organization_id=usage.organization_id,
            usage_event_id=row.id,
            amount=total,
        )
        return debit.balance_after if debit is not None else None

    async def breakdown(
        self,
        *,
        organization_id: str,
        since: datetime | None = None,
        until: datetime | None = None,
    ) -> UsageBreakdown:
        filters = [UsageEventRow.organization_id == organization_id]
        if since is not None:
            filters.append(UsageEventRow.created_at >= since)
        if until is not None:
            filters.append(UsageEventRow.created_at <= until)

        by_unit_stmt = (
            select(
                UsageEventRow.unit,
                func.coalesce(func.sum(UsageEventRow.quantity), 0.0),
                func.coalesce(func.sum(UsageEventRow.total_price), 0.0),
                func.coalesce(func.sum(UsageEventRow.internal_cost_estimate), 0.0),
            )
            .where(*filters)
            .group_by(UsageEventRow.unit)
        )
        by_unit = [
            UsageAggregate(
                unit=RateUnit(unit),
                quantity=float(qty),
                total_price=float(price),
                internal_cost_estimate=float(cost),
            )
            for unit, qty, price, cost in (await self.session.execute(by_unit_stmt)).all()
        ]

        emp_stmt = (
            select(
                UsageEventRow.employee_id,
                func.coalesce(func.sum(UsageEventRow.total_price), 0.0),
            )
            .where(*filters, UsageEventRow.employee_id.is_not(None))
            .group_by(UsageEventRow.employee_id)
        )
        by_employee = {
            str(emp): float(total)
            for emp, total in (await self.session.execute(emp_stmt)).all()
        }

        proj_stmt = (
            select(
                UsageEventRow.project_id,
                func.coalesce(func.sum(UsageEventRow.total_price), 0.0),
            )
            .where(*filters, UsageEventRow.project_id.is_not(None))
            .group_by(UsageEventRow.project_id)
        )
        by_project = {
            str(proj): float(total)
            for proj, total in (await self.session.execute(proj_stmt)).all()
        }

        return UsageBreakdown(
            by_unit=by_unit,
            by_employee=by_employee,
            by_project=by_project,
            total_price=sum(a.total_price for a in by_unit),
            total_cost=sum(a.internal_cost_estimate for a in by_unit),
        )

    async def list_events(
        self, *, organization_id: str, limit: int = 100, offset: int = 0
    ) -> list[UsageEventRow]:
        stmt = (
            select(UsageEventRow)
            .where(UsageEventRow.organization_id == organization_id)
            .order_by(UsageEventRow.created_at.desc())
            .limit(min(limit, 500))
            .offset(offset)
        )
        return list((await self.session.execute(stmt)).scalars())

    async def platform_totals(self) -> list[tuple[str, float, float]]:
        """``(organization_id, quantity, total_price)`` across every tenant.

        Platform-owner console only.
        """
        stmt = (
            select(
                UsageEventRow.organization_id,
                func.coalesce(func.sum(UsageEventRow.quantity), 0.0),
                func.coalesce(func.sum(UsageEventRow.total_price), 0.0),
            )
            .group_by(UsageEventRow.organization_id)
            .order_by(func.sum(UsageEventRow.total_price).desc())
        )
        rows = (await self.session.execute(stmt)).all()
        return [(str(org), float(qty), float(price)) for org, qty, price in rows]


class ObservabilityRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def create_trace(self, trace: Trace) -> Trace:
        self.session.add(
            TraceRow(
                id=trace.id, organization_id=trace.organization_id,
                run_id=trace.run_id, name=trace.name, created_at=trace.created_at,
            )
        )
        await self.session.flush()
        return trace

    async def add_span(self, span: Span) -> Span:
        self.session.add(
            SpanRow(
                id=span.id, trace_id=span.trace_id, parent_span_id=span.parent_span_id,
                name=span.name, kind=span.kind, attributes_json=span.attributes_json,
                started_at=span.started_at, ended_at=span.ended_at,
            )
        )
        await self.session.flush()
        return span

    async def record_error(self, event: ErrorEvent) -> None:
        self.session.add(
            ErrorEventRow(
                id=event.id,
                organization_id=event.organization_id,
                run_id=event.run_id,
                kind=event.kind,
                message=event.message,
                stack=event.stack,
                metadata_json=event.metadata_json,
                source=event.source,
                severity=event.severity,
                error_type=event.error_type,
                fingerprint=event.fingerprint,
                status_code=event.status_code,
                method=event.method,
                path=event.path,
                user_id=event.user_id,
                request_id=event.request_id,
                ip=event.ip,
                user_agent=event.user_agent,
                resolved_at=event.resolved_at,
                resolved_note=event.resolved_note,
                created_at=event.created_at,
            )
        )
        await self.session.flush()

    async def record_metric(self, metric: QualityMetric) -> None:
        self.session.add(
            QualityMetricRow(
                id=metric.id, organization_id=metric.organization_id, name=metric.name,
                value=metric.value, dimensions_json=metric.dimensions_json,
                created_at=metric.created_at,
            )
        )
        await self.session.flush()

    async def list_spans(self, *, trace_id: str) -> list[Span]:
        stmt = select(SpanRow).where(SpanRow.trace_id == trace_id).order_by(SpanRow.started_at)
        return [
            Span(
                id=r.id, trace_id=r.trace_id, parent_span_id=r.parent_span_id, name=r.name,
                kind=r.kind, attributes_json=r.attributes_json or {},
                started_at=r.started_at, ended_at=r.ended_at,
            )
            for r in (await self.session.execute(stmt)).scalars()
        ]

    # -- error console --------------------------------------------------------

    def _error_filters(
        self,
        *,
        severity: str | None = None,
        source: str | None = None,
        status: str | None = None,
        kind: str | None = None,
        organization_id: str | None = None,
        fingerprint: str | None = None,
        since: datetime | None = None,
        search: str | None = None,
    ) -> list:
        """Build the WHERE clauses shared by the list and count queries."""
        clauses: list = []
        if severity:
            clauses.append(ErrorEventRow.severity == severity)
        if source:
            clauses.append(ErrorEventRow.source == source)
        if status == "open":
            clauses.append(ErrorEventRow.resolved_at.is_(None))
        elif status == "resolved":
            clauses.append(ErrorEventRow.resolved_at.is_not(None))
        if kind:
            clauses.append(ErrorEventRow.kind == kind)
        if organization_id:
            clauses.append(ErrorEventRow.organization_id == organization_id)
        if fingerprint:
            clauses.append(ErrorEventRow.fingerprint == fingerprint)
        if since is not None:
            clauses.append(ErrorEventRow.created_at >= since)
        if search:
            clauses.append(ErrorEventRow.message.ilike(f"%{search}%"))
        return clauses

    async def list_errors(
        self,
        *,
        limit: int = 50,
        offset: int = 0,
        **filters,
    ) -> list[ErrorEvent]:
        stmt = (
            select(ErrorEventRow)
            .where(*self._error_filters(**filters))
            # Newest first: the failure someone is reacting to is almost always
            # the most recent one.
            .order_by(ErrorEventRow.created_at.desc())
            .limit(limit)
            .offset(offset)
        )
        return [_error_to_domain(row) for row in (await self.session.execute(stmt)).scalars()]

    async def count_errors(self, **filters) -> int:
        stmt = select(func.count()).select_from(ErrorEventRow).where(
            *self._error_filters(**filters)
        )
        return int((await self.session.execute(stmt)).scalar_one())

    async def get_error(self, error_id: str) -> ErrorEvent | None:
        row = (
            await self.session.execute(
                select(ErrorEventRow).where(ErrorEventRow.id == error_id)
            )
        ).scalar_one_or_none()
        return _error_to_domain(row) if row is not None else None

    async def list_error_groups(
        self,
        *,
        limit: int = 50,
        since: datetime | None = None,
        status: str | None = None,
        severity: str | None = None,
        source: str | None = None,
        organization_id: str | None = None,
    ) -> list[dict]:
        """One row per distinct bug, newest activity first.

        This is the view that makes the console usable: a bug that fired 4 000
        times is one line with a count, not 4 000 lines. ``open_count`` is what
        tells you whether anyone has actually dealt with it.
        """
        clauses: list = [ErrorEventRow.fingerprint != ""]
        if since is not None:
            clauses.append(ErrorEventRow.created_at >= since)
        if severity:
            clauses.append(ErrorEventRow.severity == severity)
        if source:
            clauses.append(ErrorEventRow.source == source)
        if organization_id:
            clauses.append(ErrorEventRow.organization_id == organization_id)

        open_count = func.sum(
            case((ErrorEventRow.resolved_at.is_(None), 1), else_=0)
        ).label("open_count")
        stmt = (
            select(
                ErrorEventRow.fingerprint,
                func.count().label("occurrences"),
                open_count,
                func.min(ErrorEventRow.created_at).label("first_seen"),
                func.max(ErrorEventRow.created_at).label("last_seen"),
                func.max(ErrorEventRow.kind).label("kind"),
                func.max(ErrorEventRow.error_type).label("error_type"),
                func.max(ErrorEventRow.severity).label("severity"),
                func.max(ErrorEventRow.source).label("source"),
                func.max(ErrorEventRow.path).label("path"),
                func.max(ErrorEventRow.message).label("sample_message"),
            )
            .where(*clauses)
            .group_by(ErrorEventRow.fingerprint)
            .order_by(func.max(ErrorEventRow.created_at).desc())
            .limit(limit)
        )
        rows = (await self.session.execute(stmt)).all()
        groups = [
            {
                "fingerprint": row.fingerprint,
                "occurrences": int(row.occurrences),
                "open_count": int(row.open_count or 0),
                "first_seen": row.first_seen,
                "last_seen": row.last_seen,
                "kind": row.kind or "",
                "error_type": row.error_type or "",
                "severity": row.severity or "",
                "source": row.source or "",
                "path": row.path or "",
                "sample_message": row.sample_message or "",
            }
            for row in rows
        ]
        if status == "open":
            return [g for g in groups if g["open_count"] > 0]
        if status == "resolved":
            return [g for g in groups if g["open_count"] == 0]
        return groups

    async def resolve_error(
        self, error_id: str, *, note: str, resolved_at: datetime
    ) -> ErrorEvent | None:
        row = (
            await self.session.execute(
                select(ErrorEventRow).where(ErrorEventRow.id == error_id)
            )
        ).scalar_one_or_none()
        if row is None:
            return None
        row.resolved_at = resolved_at
        row.resolved_note = note
        await self.session.flush()
        return _error_to_domain(row)

    async def resolve_error_group(
        self, fingerprint: str, *, note: str, resolved_at: datetime
    ) -> int:
        """Resolve every open occurrence of one bug. Returns how many changed."""
        stmt = select(ErrorEventRow).where(
            ErrorEventRow.fingerprint == fingerprint,
            ErrorEventRow.resolved_at.is_(None),
        )
        rows = list((await self.session.execute(stmt)).scalars())
        for row in rows:
            row.resolved_at = resolved_at
            row.resolved_note = note
        await self.session.flush()
        return len(rows)

    async def reopen_error_group(self, fingerprint: str) -> int:
        """Undo a resolution — for when the bug comes back."""
        stmt = select(ErrorEventRow).where(
            ErrorEventRow.fingerprint == fingerprint,
            ErrorEventRow.resolved_at.is_not(None),
        )
        rows = list((await self.session.execute(stmt)).scalars())
        for row in rows:
            row.resolved_at = None
            row.resolved_note = ""
        await self.session.flush()
        return len(rows)

    async def count_open_by_severity(self) -> dict[str, int]:
        """Unresolved counts per severity, for the console's header tiles."""
        stmt = (
            select(ErrorEventRow.severity, func.count())
            .where(ErrorEventRow.resolved_at.is_(None))
            .group_by(ErrorEventRow.severity)
        )
        return {
            str(severity): int(count)
            for severity, count in (await self.session.execute(stmt)).all()
        }