"""A unit-safe first cost model for an AI inference workload.

Every quantity and rate in the executable example is an explicit illustrative
input. No value is a vendor quote, benchmark, forecast, or recommendation.
"""

from __future__ import annotations

from dataclasses import dataclass
from decimal import Decimal

D = Decimal
GIB = D(1024**3)
SECONDS_PER_HOUR = D(3600)
TOKENS_PER_MILLION = D(1_000_000)


def _non_negative(name: str, value: Decimal) -> None:
    if not value.is_finite() or value < 0:
        raise ValueError(f"{name} must be a non-negative finite value")


@dataclass(frozen=True)
class WorkloadAssumptions:
    days: int
    requests_per_day: int
    mean_input_tokens: Decimal
    mean_output_tokens: Decimal
    accelerator_seconds_per_request: Decimal
    stored_bytes_per_request: Decimal
    egress_bytes_per_request: Decimal

    def __post_init__(self) -> None:
        if self.days <= 0 or self.requests_per_day <= 0:
            raise ValueError("days and requests_per_day must be positive")
        for name in (
            "mean_input_tokens",
            "mean_output_tokens",
            "accelerator_seconds_per_request",
            "stored_bytes_per_request",
            "egress_bytes_per_request",
        ):
            _non_negative(name, getattr(self, name))


@dataclass(frozen=True)
class UnitRates:
    input_usd_per_million_tokens: Decimal
    output_usd_per_million_tokens: Decimal
    accelerator_usd_per_hour: Decimal
    storage_usd_per_gib_month: Decimal
    egress_usd_per_gib: Decimal

    def __post_init__(self) -> None:
        for name, value in self.__dict__.items():
            _non_negative(name, value)


@dataclass(frozen=True)
class MonthlyUsage:
    requests: int
    input_tokens: Decimal
    output_tokens: Decimal
    accelerator_hours: Decimal
    end_of_period_storage_gib: Decimal
    egress_gib: Decimal


@dataclass(frozen=True)
class CostBreakdown:
    input_tokens_usd: Decimal
    output_tokens_usd: Decimal
    accelerator_usd: Decimal
    storage_usd: Decimal
    egress_usd: Decimal

    @property
    def total_usd(self) -> Decimal:
        return sum(self.__dict__.values(), start=D(0))


def monthly_usage(workload: WorkloadAssumptions) -> MonthlyUsage:
    requests = workload.days * workload.requests_per_day
    request_count = D(requests)
    return MonthlyUsage(
        requests=requests,
        input_tokens=request_count * workload.mean_input_tokens,
        output_tokens=request_count * workload.mean_output_tokens,
        accelerator_hours=(
            request_count * workload.accelerator_seconds_per_request / SECONDS_PER_HOUR
        ),
        # Deliberately simple first model: charge the end-of-period stored volume
        # for one month. Replace this with byte-hours for a real retention policy.
        end_of_period_storage_gib=(
            request_count * workload.stored_bytes_per_request / GIB
        ),
        egress_gib=request_count * workload.egress_bytes_per_request / GIB,
    )


def estimate_month(
    workload: WorkloadAssumptions,
    rates: UnitRates,
) -> tuple[MonthlyUsage, CostBreakdown]:
    usage = monthly_usage(workload)
    costs = CostBreakdown(
        input_tokens_usd=(
            usage.input_tokens / TOKENS_PER_MILLION
            * rates.input_usd_per_million_tokens
        ),
        output_tokens_usd=(
            usage.output_tokens / TOKENS_PER_MILLION
            * rates.output_usd_per_million_tokens
        ),
        accelerator_usd=usage.accelerator_hours * rates.accelerator_usd_per_hour,
        storage_usd=(
            usage.end_of_period_storage_gib * rates.storage_usd_per_gib_month
        ),
        egress_usd=usage.egress_gib * rates.egress_usd_per_gib,
    )
    return usage, costs


# These are explicit illustrative inputs, not prices, benchmarks, or forecasts.
ILLUSTRATIVE_WORKLOAD = WorkloadAssumptions(
    days=30,
    requests_per_day=10_000,
    mean_input_tokens=D("800"),
    mean_output_tokens=D("200"),
    accelerator_seconds_per_request=D("0.4"),
    stored_bytes_per_request=D("2048"),
    egress_bytes_per_request=D("4096"),
)

ILLUSTRATIVE_RATES = UnitRates(
    input_usd_per_million_tokens=D("1"),
    output_usd_per_million_tokens=D("4"),
    accelerator_usd_per_hour=D("3"),
    storage_usd_per_gib_month=D("0.10"),
    egress_usd_per_gib=D("0.05"),
)


def format_example() -> str:
    usage, costs = estimate_month(ILLUSTRATIVE_WORKLOAD, ILLUSTRATIVE_RATES)
    return "\n".join(
        [
            "example=illustrative_inputs_only",
            f"requests={usage.requests:,}",
            f"input_tokens={usage.input_tokens:,.0f}",
            f"output_tokens={usage.output_tokens:,.0f}",
            f"accelerator_hours={usage.accelerator_hours:.2f}",
            f"storage_gib={usage.end_of_period_storage_gib:.2f}",
            f"egress_gib={usage.egress_gib:.2f}",
            f"total_usd={costs.total_usd:.2f}",
            f"usd_per_request={costs.total_usd / D(usage.requests):.6f}",
        ]
    )


if __name__ == "__main__":
    print(format_example())