import psycopg
from psycopg import Connection
from src.__types.index import Tenant, UsageEventRequestMessage, UsageEventRequestPayload
from src.domain.entities.usage_event import UsageEvent
from src.infrastructure.repositories.price_repository import PriceRepository
from src.infrastructure.repositories.quota_repository import QuotaRepository
from src.infrastructure.repositories.usage_event_repository import UsageEventRepository
from src.infrastructure.usage_event_request_queue import UsageEventRequestQueue
from src.infrastructure.usage_event_request_storage import UsageEventRequestStorage
from src.logger import logger


class UsageEventUseCase:
    def __init__(
        self,
        repository: UsageEventRepository,
        quota_repository: QuotaRepository,
        price_repository: PriceRepository,
        connection: Connection,
    ) -> None:
        self.repository = repository
        self.quota_repository = quota_repository
        self.price_repository = price_repository
        self.connection = connection

    @classmethod
    def create(cls, tenant: Tenant, connection: Connection) -> "UsageEventUseCase":
        return UsageEventUseCase(
            repository=UsageEventRepository(tenant=tenant, connection=connection),
            quota_repository=QuotaRepository(tenant=tenant, connection=connection),
            price_repository=PriceRepository(connection=connection),
            connection=connection,
        )

    def validate(self, code: str) -> None:
        """使用イベントの検証を行う"""
        # 使用区分が記録されているか確認
        self.price_repository.get(code)

    def accept(
        self,
        request: UsageEventRequestPayload,
        storage: UsageEventRequestStorage,
        queue: UsageEventRequestQueue,
    ) -> None:
        """リクエスト内容を一時保存用ストレージにアップロードし、キューに登録する。

        1. リクエスト内容をストレージに一時保存する。同一箇所に既に保存済みの場合は何もしない(冪等性の保証)。
        2. アップロードしたオブジェクトのキーをキューに登録する。

        Args:
            request (UsageEventRequestPayload): 使用イベントのリクエスト内容。
            storage (UsageEventRequestStorage): 一時保存用ストレージ。
            queue (UsageEventRequestQueue): 使用イベントリクエストキュー。

        """
        object_key = storage.upload(usage_event=request)
        if not object_key:
            # Noneの場合はすでに同じオブジェクトが存在するため、キューに登録しない
            return

        queue.publish(
            message=UsageEventRequestMessage(
                tenant=request.tenant,
                object_key=object_key,
            ),
        )

    def consume(
        self,
        message: UsageEventRequestMessage,
        storage: UsageEventRequestStorage,
    ) -> None:
        """使用イベントを記録し、クォータを減算する。

        1. 一時保存した使用イベントリクエストを取得し、存在しない場合は処理は何も行わない(冪等性の保証)。
        2. idempotency_key と client_id を用いて重複リクエストを検知し、既に記録済みの場合は何も行わない(冪等性の保証)。
        3. 新規リクエストの場合、usage_events を記録する。
        4. 各 usage_event に対して、対応する単位を計算し、クォータを減算する。
        5. 一時保存用ストレージから使用イベントリクエストを削除する。

        Args:
            message (UsageEventRequestMessage): 使用イベントリクエストメッセージ。
            storage (UsageEventRequestStorage): 一時保存用ストレージ。

        """
        usage_event_request = storage.get(key=message.object_key)
        if not usage_event_request:
            # 一時保存した使用イベントリクエストが存在しない場合は何も行わない
            logger.warning(f"Usage event not found for object key: {message.object_key}. Skipping.")
            return

        try:
            self.repository.put_usage_event_request(
                client_id=usage_event_request.client.id,
                idempotency_key=usage_event_request.idempotency_key,
                timestamp=usage_event_request.timestamp,
            )
        except psycopg.errors.UniqueViolation:
            # 重複している場合は何もしない
            logger.debug(f"Duplicate usage event request with idempotency key: {usage_event_request.idempotency_key}")
            return

        usage_events = [
            UsageEvent.model_validate(request_event) for request_event in usage_event_request.request_events
        ]
        for usage_event in usage_events:
            logger.debug(f"Recording usage event: {usage_event}")
            unit = usage_event.calc_unit(
                self.price_repository.get(usage_event.code),
            )
            # 使用イベントをデータベースに記録
            # クォータ割り当て処理とは独立したトランザクションで実行し、記録の確実性を向上
            # この段階ではユニット使用量の記録のみで、クォータから差し引く処理は後続のトランザクションで行う
            usage_event_id = self.repository.put(
                client_id=usage_event_request.client.id,
                client_tenant_id=usage_event_request.client_tenant.id,
                usage_event=usage_event,
                unit=unit,
            )

            # ユニット使用量をクォータ割り当てから差し引く処理
            # 現在使用中のクォータ割り当て残量がユニット使用量に満たない場合は、可能な分のみユニット使用量として記録する
            # その後に、次に優先度の高いクォータ割り当てからユニット使用量の残りを差し引く
            # 全てのユニット使用量を差し引くまで繰り返し処理を行う
            while unit != 0:
                with self.connection.transaction():
                    quota = self.quota_repository.get()
                    # クォータが超過している場合は、クォータ割り当てからの差し引きを行わずにユニット使用量を記録
                    if quota.is_exceeded:
                        self.repository.put_quota_usage(
                            usage_event_id=usage_event_id,
                            quota_allocation=None,
                            unit=unit,
                        )
                        break
                    if quota.in_use_allocation_id is None:
                        break
                    # 使用中のクォータ割り当てからユニット使用量を差し引くため、ロックする
                    in_use_quota_allocation = self.quota_repository.lock_quota_allocation(
                        allocation_id=quota.in_use_allocation_id,
                    )
                    remaining_unit = in_use_quota_allocation.use(unit)
                    # どのクォータ割り当てからユニット使用量を差し引いたかを記録
                    self.repository.put_quota_usage(
                        usage_event_id=usage_event_id,
                        quota_allocation=in_use_quota_allocation,
                        unit=unit - remaining_unit,
                    )
                unit = remaining_unit

        # 正常終了時はストレージからリクエストを削除
        storage.delete(key=message.object_key)
