amazon-transcreation/backend/app/services/report_service.py

from datetime import datetime
from typing import Any
from uuid import UUID

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

from app.models.audit import TokenUsageLog
from app.models.feedback import Feedback, FlagType
from app.models.job import Job, JobStatus, LocaleInstance, LocaleStatus
from app.models.output import OutputRow, ConfidenceTier


class ReportService:
    """Service for aggregation queries powering reports."""

    async def get_usage_stats(
        self,
        db: AsyncSession,
        client_id: UUID | None = None,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
    ) -> dict[str, Any]:
        """Get overall usage statistics."""
        job_query = select(
            func.count(Job.id).label("total_jobs"),
            func.sum(Job.total_token_usage).label("total_tokens"),
            func.sum(Job.total_estimated_cost).label("total_cost"),
        )
        if client_id:
            job_query = job_query.where(Job.client_id == client_id)
        if date_from:
            job_query = job_query.where(Job.created_at >= date_from)
        if date_to:
            job_query = job_query.where(Job.created_at <= date_to)

        result = await db.execute(job_query)
        row = result.one()

        # Status breakdown (apply same filters as main query)
        status_query = select(
            Job.status, func.count(Job.id)
        ).group_by(Job.status)
        if client_id:
            status_query = status_query.where(Job.client_id == client_id)
        if date_from:
            status_query = status_query.where(Job.created_at >= date_from)
        if date_to:
            status_query = status_query.where(Job.created_at <= date_to)
        status_result = await db.execute(status_query)
        status_breakdown = {
            status.value: count for status, count in status_result.all()
        }

        return {
            "total_jobs": row.total_jobs or 0,
            "total_tokens": row.total_tokens or 0,
            "total_cost": float(row.total_cost or 0.0),
            "status_breakdown": status_breakdown,
        }

    async def get_token_cost_data(
        self,
        db: AsyncSession,
        client_id: UUID | None = None,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
    ) -> list[dict[str, Any]]:
        """Get token usage and cost data grouped by date."""
        day_expr = func.date_trunc(literal_column("'day'"), TokenUsageLog.timestamp)
        query = (
            select(
                day_expr.label("date"),
                func.sum(TokenUsageLog.input_tokens).label("input_tokens"),
                func.sum(TokenUsageLog.output_tokens).label("output_tokens"),
                func.sum(TokenUsageLog.total_tokens).label("total_tokens"),
                func.sum(TokenUsageLog.estimated_cost_usd).label("total_cost"),
            )
            .group_by(day_expr)
            .order_by(day_expr)
        )

        if client_id:
            query = query.join(LocaleInstance).join(Job).where(
                Job.client_id == client_id
            )
        if date_from:
            query = query.where(TokenUsageLog.timestamp >= date_from)
        if date_to:
            query = query.where(TokenUsageLog.timestamp <= date_to)

        result = await db.execute(query)
        return [
            {
                "date": str(row.date),
                "input_tokens": row.input_tokens or 0,
                "output_tokens": row.output_tokens or 0,
                "total_tokens": row.total_tokens or 0,
                "total_cost": float(row.total_cost or 0.0),
            }
            for row in result.all()
        ]

    async def get_quality_metrics(
        self,
        db: AsyncSession,
        client_id: UUID | None = None,
    ) -> dict[str, Any]:
        """Get quality metrics from output confidence tiers and feedback."""
        # Confidence tier distribution
        tier_query = select(
            OutputRow.confidence_tier, func.count(OutputRow.id)
        ).group_by(OutputRow.confidence_tier)

        if client_id:
            tier_query = tier_query.join(LocaleInstance).join(Job).where(
                Job.client_id == client_id
            )

        tier_result = await db.execute(tier_query)
        tier_breakdown = {
            tier.value: count for tier, count in tier_result.all()
        }

        # Feedback distribution
        feedback_query = select(
            Feedback.flag_type, func.count(Feedback.id)
        ).group_by(Feedback.flag_type)

        feedback_result = await db.execute(feedback_query)
        feedback_breakdown = {
            ft.value: count for ft, count in feedback_result.all()
        }

        return {
            "confidence_tiers": tier_breakdown,
            "feedback_distribution": feedback_breakdown,
        }

    async def get_locale_stats(
        self,
        db: AsyncSession,
        client_id: UUID | None = None,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
    ) -> list[dict[str, Any]]:
        """Get per-locale aggregate stats (tokens, cost, avg duration, job count)."""
        query = select(
            LocaleInstance.locale_code,
            func.count(LocaleInstance.id).label("count"),
            func.sum(LocaleInstance.token_usage).label("total_tokens"),
            func.sum(LocaleInstance.estimated_cost).label("total_cost"),
            func.avg(
                func.extract("epoch", LocaleInstance.completed_at)
                - func.extract("epoch", LocaleInstance.started_at)
            ).label("avg_duration_seconds"),
        ).where(
            LocaleInstance.status == LocaleStatus.complete,
            LocaleInstance.started_at.isnot(None),
            LocaleInstance.completed_at.isnot(None),
        ).group_by(LocaleInstance.locale_code)

        needs_job_join = bool(client_id or date_from or date_to)
        if needs_job_join:
            query = query.join(Job)
        if client_id:
            query = query.where(Job.client_id == client_id)
        if date_from:
            query = query.where(Job.created_at >= date_from)
        if date_to:
            query = query.where(Job.created_at <= date_to)

        result = await db.execute(query)
        out: list[dict[str, Any]] = []
        for row in result.all():
            count = row.count or 0
            total_cost = float(row.total_cost or 0.0)
            out.append({
                "locale": row.locale_code,
                "count": count,
                "total_tokens": row.total_tokens or 0,
                "total_cost": total_cost,
                "avg_cost_per_brief": round(total_cost / count, 4) if count else 0.0,
                "avg_duration_minutes": round(float(row.avg_duration_seconds or 0) / 60, 1),
            })
        return out

    async def get_jobs_over_time(
        self,
        db: AsyncSession,
        client_id: UUID | None = None,
        date_from: datetime | None = None,
        date_to: datetime | None = None,
    ) -> list[dict[str, Any]]:
        """Get job counts grouped by week."""
        week_expr = func.date_trunc(literal_column("'week'"), Job.created_at)

        query = select(
            week_expr.label("week"),
            func.count(Job.id).label("total"),
            func.sum(case((Job.status == JobStatus.complete, 1), else_=0)).label("completed_raw"),
            func.sum(case((Job.status == JobStatus.error, 1), else_=0)).label("errors_raw"),
        ).group_by(
            week_expr
        ).order_by(
            week_expr
        )

        if client_id:
            query = query.where(Job.client_id == client_id)
        if date_from:
            query = query.where(Job.created_at >= date_from)
        if date_to:
            query = query.where(Job.created_at <= date_to)

        result = await db.execute(query)
        rows = []
        for row in result.all():
            week_str = row.week.strftime("%b %d") if row.week else "?"
            rows.append({
                "name": week_str,
                "jobs": row.total or 0,
                "completed": int(row.completed_raw or 0),
                "errors": int(row.errors_raw or 0),
            })
        return rows