from decimal import Decimal

from psycopg import Connection
from psycopg.rows import class_row
from src.__types.index import GroupBy, QuotaUsage, Tenant, Unit
from src.domain.entities.quota_allocation import QuotaAllocationStatus
from src.infrastructure.repositories.quota_usage_aggregate_query_builder import QuotaUsageAggregateQueryBuilder
from src.tracer import tracer


class QuotaUsageRepository:
    """クォータ使用実績を管理するリポジトリ"""

    def __init__(self, tenant: Tenant, connection: Connection) -> None:
        self.tenant = tenant
        self.connection = connection

    @tracer.capture_method
    def get_last_month_usage(self) -> Unit:
        """前月の使用量を取得する"""
        with self.connection.cursor() as cur:
            cur.execute(
                """--sql
                WITH cte_usage_events AS (
                    SELECT id
                    FROM usage_events
                    WHERE tenant_id = %(tenant_id)s
                    AND timestamp >= date_trunc('month', current_date - interval '1 month')
                    AND timestamp < date_trunc('month', current_date)
                )
                SELECT SUM(unit)
                FROM quota_usages INNER JOIN cte_usage_events ON quota_usages.usage_event_id = cte_usage_events.id
                WHERE tenant_id = %(tenant_id)s
                AND quota_allocation_id IS NOT NULL -- クォータ割り当てに紐づけられなかったユニット量は除外
                """,
                {"tenant_id": self.tenant.id},
            )
            if row := cur.fetchone():
                if row[0] is None:
                    return Decimal(0)
                return row[0]

        return Decimal(0)

    @tracer.capture_method
    def aggregate(
        self,
        group_by: list[GroupBy],
        quota_allocation_status: list[QuotaAllocationStatus],
        year: int | None = None,
    ) -> list[QuotaUsage]:
        """クォータ使用実績を集計する"""
        query_builder = QuotaUsageAggregateQueryBuilder(
            tenant_id=self.tenant.id,
            status=quota_allocation_status,
            group_by=group_by,
            year=year,
        )
        query, parameters = query_builder.build()
        with self.connection.cursor(row_factory=class_row(QuotaUsage)) as cur:
            cur.execute(query, parameters)
            return cur.fetchall()
