from datetime import datetime
from decimal import Decimal
from uuid import UUID

import psycopg
from src.__types.index import Tenant, Unit
from src.domain.entities.quota_allocation import QuotaAllocation
from src.domain.entities.usage_event import UsageEvent


class UsageEventRepository:
    """使用イベントを管理するリポジトリ"""

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

    def put_quota_usage(
        self,
        usage_event_id: UUID,
        quota_allocation: QuotaAllocation | None,
        unit: Decimal,
    ) -> None:
        """使用量をどのクォータ割り当てから差し引いたかを記録する"""
        with self.connection.cursor() as cur:
            cur.execute(
                """--sql
                INSERT INTO quota_usages (tenant_id, usage_event_id, quota_allocation_id, unit)
                VALUES (%(tenant_id)s, %(usage_event_id)s, %(quota_allocation_id)s, %(unit)s);
                """,
                {
                    "tenant_id": self.tenant.id,
                    "usage_event_id": usage_event_id,
                    "quota_allocation_id": quota_allocation.id if quota_allocation else None,
                    "unit": unit,
                },
            )
            if quota_allocation:
                cur.execute(
                    """--sql
                    UPDATE quota_allocations
                    SET used = %(used)s, status = %(status)s
                    WHERE id = %(quota_allocation_id)s
                    AND tenant_id = %(tenant_id)s
                    """,
                    {
                        "tenant_id": self.tenant.id,
                        "used": quota_allocation.used,
                        "quota_allocation_id": quota_allocation.id,
                        "status": quota_allocation.status.value,
                    },
                )

    def put_usage_event_request(
        self,
        client_id: UUID,
        idempotency_key: UUID,
        timestamp: datetime,
    ) -> None:
        """使用イベントリクエストを記録する"""
        with self.connection.cursor() as cur:
            cur.execute(
                """--sql
                INSERT INTO usage_event_requests (client_id, idempotency_key, "timestamp")
                VALUES (%(client_id)s, %(idempotency_key)s, %(timestamp)s)
                """,
                {
                    "client_id": client_id,
                    "idempotency_key": idempotency_key,
                    "timestamp": timestamp,
                },
            )

    def put(
        self,
        client_id: UUID,
        client_tenant_id: UUID,
        usage_event: UsageEvent,
        unit: Unit,
    ) -> UUID:
        """使用イベントを記録する"""
        with self.connection.cursor() as cur:
            cur.execute(
                """--sql
                INSERT INTO usage_events (
                    tenant_id,
                    client_tenant_id,
                    client_id,
                    code,
                    quantity,
                    unit,
                    service,
                    feature,
                    user_id,
                    trace_id
                ) VALUES (
                    %(tenant_id)s,
                    %(client_tenant_id)s,
                    %(client_id)s,
                    %(code)s,
                    %(quantity)s,
                    %(unit)s,
                    %(service)s,
                    %(feature)s,
                    %(user_id)s,
                    %(trace_id)s
                )
                RETURNING id
                """,
                {
                    "tenant_id": self.tenant.id,
                    "client_tenant_id": client_tenant_id,
                    "client_id": client_id,
                    "code": usage_event.code,
                    "quantity": usage_event.quantity,
                    "unit": unit,
                    "service": usage_event.service,
                    "feature": usage_event.feature,
                    "user_id": usage_event.user_id,
                    "trace_id": usage_event.trace_id,
                },
            )
            row = cur.fetchone()
            if not row:
                raise ValueError("Failed to get inserted usage_event id.")
            return row[0]  # id
