vault-dash/app/domain/units.py

from __future__ import annotations

from dataclasses import dataclass
from decimal import Decimal
from enum import StrEnum


class BaseCurrency(StrEnum):
    USD = "USD"
    EUR = "EUR"
    CHF = "CHF"


class WeightUnit(StrEnum):
    GRAM = "g"
    KILOGRAM = "kg"
    OUNCE_TROY = "ozt"


GRAMS_PER_KILOGRAM = Decimal("1000")
GRAMS_PER_TROY_OUNCE = Decimal("31.1034768")


DecimalLike = Decimal | int | str


def to_decimal(value: DecimalLike) -> Decimal:
    if isinstance(value, bool):
        raise TypeError("Boolean values are not valid Decimal inputs")
    if isinstance(value, Decimal):
        amount = value
    elif isinstance(value, int):
        amount = Decimal(value)
    elif isinstance(value, str):
        amount = Decimal(value)
    else:
        raise TypeError(f"Unsupported decimal input type: {type(value)!r}")
    if not amount.is_finite():
        raise ValueError("Decimal value must be finite")
    return amount


def decimal_from_float(value: float) -> Decimal:
    if not isinstance(value, float):
        raise TypeError(f"Expected float, got {type(value)!r}")
    amount = Decimal(str(value))
    if not amount.is_finite():
        raise ValueError("Decimal value must be finite")
    return amount


def _coerce_currency(value: BaseCurrency | str) -> BaseCurrency:
    if isinstance(value, BaseCurrency):
        return value
    try:
        return BaseCurrency(value)
    except ValueError as exc:
        raise ValueError(f"Invalid currency: {value!r}") from exc


def _coerce_weight_unit(value: WeightUnit | str) -> WeightUnit:
    if isinstance(value, WeightUnit):
        return value
    try:
        return WeightUnit(value)
    except ValueError as exc:
        raise ValueError(f"Invalid weight unit: {value!r}") from exc


def weight_unit_factor(unit: WeightUnit) -> Decimal:
    if unit is WeightUnit.GRAM:
        return Decimal("1")
    if unit is WeightUnit.KILOGRAM:
        return GRAMS_PER_KILOGRAM
    if unit is WeightUnit.OUNCE_TROY:
        return GRAMS_PER_TROY_OUNCE
    raise ValueError(f"Unsupported weight unit: {unit}")


def convert_weight(amount: Decimal, from_unit: WeightUnit, to_unit: WeightUnit) -> Decimal:
    if from_unit is to_unit:
        return amount
    grams = amount * weight_unit_factor(from_unit)
    return grams / weight_unit_factor(to_unit)


def convert_price_per_weight(amount: Decimal, from_unit: WeightUnit, to_unit: WeightUnit) -> Decimal:
    if from_unit is to_unit:
        return amount
    return amount * weight_unit_factor(to_unit) / weight_unit_factor(from_unit)


@dataclass(frozen=True, slots=True)
class Money:
    amount: Decimal
    currency: BaseCurrency | str

    def __post_init__(self) -> None:
        object.__setattr__(self, "amount", to_decimal(self.amount))
        object.__setattr__(self, "currency", _coerce_currency(self.currency))

    @classmethod
    def zero(cls, currency: BaseCurrency) -> Money:
        return cls(amount=Decimal("0"), currency=currency)

    def assert_currency(self, currency: BaseCurrency) -> Money:
        if self.currency is not currency:
            raise ValueError(f"Currency mismatch: {self.currency} != {currency}")
        return self

    def __add__(self, other: object) -> Money:
        if not isinstance(other, Money):
            return NotImplemented
        if self.currency is not other.currency:
            raise ValueError(f"Currency mismatch: {self.currency} != {other.currency}")
        return Money(amount=self.amount + other.amount, currency=self.currency)

    def __sub__(self, other: object) -> Money:
        if not isinstance(other, Money):
            return NotImplemented
        if self.currency is not other.currency:
            raise ValueError(f"Currency mismatch: {self.currency} != {other.currency}")
        return Money(amount=self.amount - other.amount, currency=self.currency)

    def __mul__(self, other: object) -> Money:
        if isinstance(other, bool):
            return NotImplemented
        if isinstance(other, Decimal):
            return Money(amount=self.amount * other, currency=self.currency)
        if isinstance(other, int):
            return Money(amount=self.amount * Decimal(other), currency=self.currency)
        return NotImplemented

    def __rmul__(self, other: object) -> Money:
        return self.__mul__(other)

    def __truediv__(self, other: object) -> Money:
        if isinstance(other, bool):
            return NotImplemented
        if isinstance(other, Decimal):
            return Money(amount=self.amount / other, currency=self.currency)
        if isinstance(other, int):
            return Money(amount=self.amount / Decimal(other), currency=self.currency)
        return NotImplemented

    def __neg__(self) -> Money:
        return Money(amount=-self.amount, currency=self.currency)


@dataclass(frozen=True, slots=True)
class Weight:
    amount: Decimal
    unit: WeightUnit | str

    def __post_init__(self) -> None:
        object.__setattr__(self, "amount", to_decimal(self.amount))
        object.__setattr__(self, "unit", _coerce_weight_unit(self.unit))

    def to_unit(self, unit: WeightUnit) -> Weight:
        return Weight(amount=convert_weight(self.amount, self.unit, unit), unit=unit)

    def __add__(self, other: object) -> Weight:
        if not isinstance(other, Weight):
            return NotImplemented
        other_converted = other.to_unit(self.unit)
        return Weight(amount=self.amount + other_converted.amount, unit=self.unit)

    def __sub__(self, other: object) -> Weight:
        if not isinstance(other, Weight):
            return NotImplemented
        other_converted = other.to_unit(self.unit)
        return Weight(amount=self.amount - other_converted.amount, unit=self.unit)

    def __mul__(self, other: object) -> Weight | Money:
        if isinstance(other, bool):
            return NotImplemented
        if isinstance(other, Decimal):
            return Weight(amount=self.amount * other, unit=self.unit)
        if isinstance(other, int):
            return Weight(amount=self.amount * Decimal(other), unit=self.unit)
        if isinstance(other, PricePerWeight):
            adjusted_weight = self.to_unit(other.per_unit)
            return Money(amount=adjusted_weight.amount * other.amount, currency=other.currency)
        return NotImplemented

    def __rmul__(self, other: object) -> Weight:
        if isinstance(other, bool):
            return NotImplemented
        if isinstance(other, Decimal):
            return Weight(amount=self.amount * other, unit=self.unit)
        if isinstance(other, int):
            return Weight(amount=self.amount * Decimal(other), unit=self.unit)
        return NotImplemented

    def __truediv__(self, other: object) -> Weight:
        if isinstance(other, bool):
            return NotImplemented
        if isinstance(other, Decimal):
            return Weight(amount=self.amount / other, unit=self.unit)
        if isinstance(other, int):
            return Weight(amount=self.amount / Decimal(other), unit=self.unit)
        return NotImplemented


@dataclass(frozen=True, slots=True)
class PricePerWeight:
    amount: Decimal
    currency: BaseCurrency | str
    per_unit: WeightUnit | str

    def __post_init__(self) -> None:
        amount = to_decimal(self.amount)
        if amount < 0:
            raise ValueError("PricePerWeight amount must be non-negative")
        object.__setattr__(self, "amount", amount)
        object.__setattr__(self, "currency", _coerce_currency(self.currency))
        object.__setattr__(self, "per_unit", _coerce_weight_unit(self.per_unit))

    def to_unit(self, unit: WeightUnit) -> PricePerWeight:
        return PricePerWeight(
            amount=convert_price_per_weight(self.amount, self.per_unit, unit),
            currency=self.currency,
            per_unit=unit,
        )

    def __mul__(self, other: object) -> Money | PricePerWeight:
        if isinstance(other, bool):
            return NotImplemented
        if isinstance(other, Weight):
            adjusted_weight = other.to_unit(self.per_unit)
            return Money(amount=adjusted_weight.amount * self.amount, currency=self.currency)
        if isinstance(other, Decimal):
            return PricePerWeight(amount=self.amount * other, currency=self.currency, per_unit=self.per_unit)
        if isinstance(other, int):
            return PricePerWeight(amount=self.amount * Decimal(other), currency=self.currency, per_unit=self.per_unit)
        return NotImplemented

    def __rmul__(self, other: object) -> PricePerWeight:
        if isinstance(other, bool):
            return NotImplemented
        if isinstance(other, Decimal):
            return PricePerWeight(amount=self.amount * other, currency=self.currency, per_unit=self.per_unit)
        if isinstance(other, int):
            return PricePerWeight(amount=self.amount * Decimal(other), currency=self.currency, per_unit=self.per_unit)
        return NotImplemented