src/services/usage_tracking.py

import logging
from datetime import datetime

from sqlalchemy import func, select
from sqlalchemy.orm import Session

from src.services.db import UsageEvent, get_session, new_id

logger = logging.getLogger(__name__)


class UsageTrackingService:
    """Usage event recording and querying.

    CRITICAL DESIGN RULE:
    - record_event_safe() opens its OWN short-lived DB session.
      It never commits, rollbacks, or closes the caller's request-scoped session.
      If the insert fails, it logs the error and returns silently — the main
      user request is never broken by a tracking failure.
    - Counting/query methods (count_user_events, count_user_events_grouped)
      use the injected request session (self.db) for read-only queries. This is
      safe because they do not modify state.
    """

    def __init__(self, db: Session):
        self.db = db  # request-scoped session, used for reads only

    @staticmethod
    def record_event_safe(
        *,
        event_type: str,
        user_id: str | None = None,
        organization_id: str | None = None,
        guest_session_id: str | None = None,
        source: str | None = None,
        metadata: dict | None = None,
        status: str = "success",
    ) -> None:
        """Insert a usage event in an INDEPENDENT session.

        Opens a new short-lived Session, inserts, commits, closes.
        On any failure: logs the error, does NOT raise, does NOT touch
        the caller's transaction.
        """
        try:
            gen = get_session()
            session = next(gen)
            try:
                event = UsageEvent(
                    id=new_id("evt"),
                    user_id=user_id,
                    organization_id=organization_id,
                    guest_session_id=guest_session_id,
                    event_type=event_type,
                    status=status,
                    source=source,
                    metadata_json=metadata,
                )
                session.add(event)
                session.commit()
            except Exception:
                session.rollback()
                raise
            finally:
                session.close()
        except Exception:
            logger.exception("Usage tracking failed (event_type=%s)", event_type)

    def count_user_events(
        self,
        user_id: str,
        event_type: str,
        *,
        start_at: datetime | None = None,
        end_before: datetime | None = None,
    ) -> int:
        """Count successful events for a single user and type."""
        stmt = select(func.count(UsageEvent.id)).where(
            UsageEvent.user_id == user_id,
            UsageEvent.event_type == event_type,
            UsageEvent.status == "success",
        )
        if start_at is not None:
            stmt = stmt.where(UsageEvent.created_at >= start_at)
        if end_before is not None:
            stmt = stmt.where(UsageEvent.created_at < end_before)
        return int(self.db.scalar(stmt) or 0)

    def count_user_events_grouped(
        self,
        user_id: str,
        event_types: list[str],
        *,
        start_at: datetime | None = None,
        end_before: datetime | None = None,
    ) -> dict[str, int]:
        """Count events grouped by event_type for a single user.
        Only counts status == 'success'.
        """
        stmt = (
            select(UsageEvent.event_type, func.count(UsageEvent.id))
            .where(
                UsageEvent.user_id == user_id,
                UsageEvent.event_type.in_(event_types),
                UsageEvent.status == "success",
            )
            .group_by(UsageEvent.event_type)
        )
        if start_at is not None:
            stmt = stmt.where(UsageEvent.created_at >= start_at)
        if end_before is not None:
            stmt = stmt.where(UsageEvent.created_at < end_before)
        rows = self.db.execute(stmt).all()
        result = {et: 0 for et in event_types}
        for event_type_val, count in rows:
            result[event_type_val] = int(count)
        return result