from decimal import Decimal

from pydantic import BaseModel, Field, validate_call
from src.__types.index import Price, Unit
from src.const import UNIT_PRICE


class UsageEvent(BaseModel):
    """使用イベント

    使用イベントとは、『クライアントから送信された生成AIの使用量』を意味する。
    AIユニットの使用量を計算、及び集計をするために必要な情報を含む。
    データベースに格納された使用区分ごとの単価情報から、AIユニット使用量を計算する。
    AIユニット使用量は以下の計算式で算出される。

    AIユニット量 = 使用量 * 使用区分毎の単価 * AIユニット単価

    使用区分はクライアントの生成AIの使用量を金額に換算するためのキーであり、例えば以下のようなものがある。
    - Claude 3 Haikuの入力トークン
    - Claude 3 Haikuの出力トークン

    使用量は使用区分毎に単位が異なる。トークン数や呼び出し回数などが想定される。
    """

    code: str = Field(..., description="使用区分")
    quantity: Decimal = Field(..., description="使用量")
    service: str = Field(..., description="サービス名")
    feature: str = Field(..., description="機能名")
    user_id: str | None = Field(default=None, description="ユーザーID")
    trace_id: str | None = Field(default=None, description="トレースID")

    @validate_call(validate_return=True)  # Unit型エイリアスのValidatorを起動させるため
    def calc_unit(self, price: Price) -> Unit:
        """ユニット使用量を計算する ※ この関数以外でAIユニット使用量を計算しないこと"""
        cost = self.quantity * price.price
        return cost / UNIT_PRICE
