"""Compute accounting helpers for inference workloads."""

from __future__ import annotations

from decimal import Decimal

from rc.credits import CREDIT_SUBUNITS_PER_CREDIT

FLOPS_PER_MFLOP = 1_000_000
MFLOPS_PER_GFLOP = 1_000
GFLOPS_PER_TFLOP = 1_000
MFLOPS_PER_TFLOP = FLOPS_PER_MFLOP
FLOPS_PER_TFLOP = 1_000_000_000_000
FLOPS_PER_CREDIT_SUBUNIT = FLOPS_PER_TFLOP // CREDIT_SUBUNITS_PER_CREDIT


def compute_flops(
    model_params: int,
    prompt_tokens: int,
    generated_tokens: int,
) -> int:
    """Return exact token-accounted FLOPs for a generation run."""
    total_tokens = max(int(prompt_tokens), 0) + max(int(generated_tokens), 0)
    return 2 * max(int(model_params), 0) * total_tokens


def compute_mflops(
    model_params: int,
    prompt_tokens: int,
    generated_tokens: int,
) -> int:
    """Return floored mega-FLOPs for compatibility with legacy telemetry."""
    return compute_flops(
        model_params=model_params,
        prompt_tokens=prompt_tokens,
        generated_tokens=generated_tokens,
    ) // FLOPS_PER_MFLOP


def flops_to_tflops_decimal(value: int) -> Decimal:
    """Convert exact raw FLOPs into an exact TFLOP decimal."""
    return Decimal(max(int(value), 0)) / Decimal(FLOPS_PER_TFLOP)


def compute_tflops(
    model_params: int,
    prompt_tokens: int,
    generated_tokens: int,
) -> float:
    """Return tera-FLOPs for a generation run as a float view."""
    return float(
        flops_to_tflops_decimal(
            compute_flops(
                model_params=model_params,
                prompt_tokens=prompt_tokens,
                generated_tokens=generated_tokens,
            )
        )
    )


def mflops_to_tflops(value: int) -> float:
    """Convert raw MFLOP telemetry into a display-friendly TFLOP value."""
    return max(int(value), 0) / MFLOPS_PER_TFLOP


def flops_to_tflops(value: int) -> float:
    """Convert exact raw FLOPs into a display-friendly TFLOP value."""
    return float(flops_to_tflops_decimal(value))


def format_flops_as_tflops(value: int) -> str:
    """Format exact raw FLOPs as an exact TFLOP decimal string."""
    text = format(flops_to_tflops_decimal(value), "f")
    return text.rstrip("0").rstrip(".") or "0"


def compute_credit_subunits(
    model_params: int,
    prompt_tokens: int,
    generated_tokens: int,
) -> int:
    """Return floored billable credit subunits for a generation run."""
    return compute_flops(
        model_params=model_params,
        prompt_tokens=prompt_tokens,
        generated_tokens=generated_tokens,
    ) // FLOPS_PER_CREDIT_SUBUNIT


def credit_subunits_to_tflops(value: int) -> float:
    """Convert raw credit subunits back into TFLOP credits."""
    return max(int(value), 0) / CREDIT_SUBUNITS_PER_CREDIT
