Source code for ferrosoft.apps.emiflow.services.tonnage

#  Copyright (c) 2025 Ferrosoft GmbH. All rights reserved.
"""Tonnage accounting for subscription billing.

Customers are billed on the cumulative weight of (non-simulation)
emission bookings. This module exposes two query paths:

* A direct database aggregate per tenant (:meth:`TonnageService.get_summary`)
  used at startup to seed the Prometheus counter and on demand for
  reconciliation. Expensive — it opens a connection to every active
  tenant database.
* A cheap VictoriaMetrics query (:meth:`TonnageServiceVictoria.get_org_stats`)
  that returns the rolling 4-week tonnage and configured limit per
  organisation from the metrics store.
"""
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass
from decimal import Decimal
from urllib.parse import urljoin

import requests
from django.core.exceptions import ImproperlyConfigured
from django.db.models.aggregates import Sum

from ferrosoft.apps.emiflow.models import EmissionBookingRecord
from ferrosoft.apps.ferrobase.models import Organization, Tenant, MeasurementUnit
from ferrosoft.apps.ferrobase.models.decimal import DecimalWithUnit
from ferrosoft.apps.ferrobase.tenant.context import active_tenant
from ferrosoft.config import require_settings
from ferrosoft.decimal import zero




@dataclass
class TonnageSummary:
    """Aggregated booked tonnage for one ``(organization, tenant)`` pair."""

    organization: Organization
    tenant: Tenant
    tonnage: Decimal





@dataclass
class TonnageStats:
    """Billing snapshot for an organisation: current usage versus configured limit."""

    billable_tonnage: DecimalWithUnit
    tonnage_limit: DecimalWithUnit





class TonnageService(ABC):
    """Abstract tonnage service. The concrete singleton is
    :class:`TonnageServiceVictoria`."""



    @classmethod
    def get_instance(cls) -> "TonnageService":
        """Return the configured production tonnage service."""
        return TonnageServiceVictoria.get_instance()




    def get_summary(self) -> list[TonnageSummary]:
        """Return per-tenant booked tonnage across every active tenant.

        Iterates over every organisation and active tenant, opening
        each tenant database to read its booking-record aggregate.
        Misconfigured databases (``ImproperlyConfigured``) yield a
        zero-tonnage row so the summary remains complete. Costly;
        prefer :meth:`get_org_stats` whenever per-organisation totals
        are sufficient.
        """
        summaries = []

        for org in Organization.objects.all():
            for tenant in org.tenants.filter(active=True):
                try:
                    tonnage = self.get_tenant_tonnage(tenant)
                    summaries.append(TonnageSummary(org, tenant, tonnage))
                except ImproperlyConfigured:
                    # This case may happen with false database credentials.
                    summaries.append(TonnageSummary(org, tenant, zero))

        return summaries




    @staticmethod
    def get_tenant_tonnage(tenant):
        """Return ``tenant``'s booked tonnage from its own database.

        Simulation records are excluded so dry-run bookings don't count
        toward the customer's subscription quota.
        """
        with active_tenant(tenant):
            records = EmissionBookingRecord.objects.filter(simulation=False)
            records = records.aggregate(Sum("weight_value", default=zero))
            return records["weight_value__sum"]




    @abstractmethod
    def get_org_stats(self, org: Organization) -> TonnageStats:
        """Return the billable tonnage and limit for ``org`` (implementation-specific)."""
        raise NotImplementedError()






class TonnageServiceVictoria(TonnageService):
    """Tonnage service backed by VictoriaMetrics Prometheus-compatible queries.

    Uses a single PromQL ``union`` over the ``emiflow_tonnage_limit``
    gauge and the 4-week increase of the ``emiflow_tonnage_total``
    counter to return both values in one HTTP call. Returns
    zero-stats on query failure so callers can render a safe fallback.
    """

    _METRIC_TON = MeasurementUnit(symbol="t")
    _ZERO_TON = DecimalWithUnit(zero, _METRIC_TON)
    _ZERO_STATS = TonnageStats(_ZERO_TON, _ZERO_TON)

    _STATS_QUERY = """union(
      alias(emiflow_tonnage_limit{organization="%(org)s"}, "tonnage_limit"),
      alias(
        sum(increase(emiflow_tonnage_total{organization="%(org)s"}[4w])) by (organization),
        "current_tonnage"
      ))"""



    @dataclass
    class Config:
        """VictoriaMetrics endpoint configuration."""

        base_url: str
        timeout: float = 10



        @classmethod
        def from_django_settings(cls):
            """Build a config from the ``VICTORIA_BASE_URL`` Django setting."""
            settings = require_settings("VICTORIA_BASE_URL")
            return cls(settings["VICTORIA_BASE_URL"])





    @classmethod
    def get_instance(cls):
        """Instantiate the service with config loaded from Django settings."""
        return cls(cls.Config.from_django_settings())


    def __init__(self, config: Config):
        self.config = config



    def get_org_stats(self, org: Organization) -> TonnageStats:
        """Return ``org``'s current tonnage and limit from VictoriaMetrics.

        Falls back to a zero stats tuple (and logs the response) when
        the query fails, returns a non-``success`` status, or yields an
        unexpected result type.
        """
        org_name = self._escape_string(org.name)
        response = requests.post(
            urljoin(self.config.base_url, "api/v1/query"),
            data={
                "query": self._STATS_QUERY % {"org": org_name},
            },
            timeout=self.config.timeout,
        )
        query_result = response.json()
        if query_result.get("status", None) != "success":
            logging.error("Victoria query failed: %s", query_result)
            return self._ZERO_STATS

        data = query_result.get("data", {})
        result_type = data.get("resultType", None)
        if result_type != "vector":
            logging.error("Victoria query returned unexpected result: %s", query_result)
            return self._ZERO_STATS

        result = data.get("result", [])
        if len(result) == 0:
            return self._ZERO_STATS

        values = self._extract_metric_values(result)

        return TonnageStats(
            DecimalWithUnit(values.get("current_tonnage", zero), self._METRIC_TON),
            DecimalWithUnit(values.get("tonnage_limit", zero), self._METRIC_TON),
        )


    @staticmethod
    def _escape_string(value: str) -> str:
        """Escape double quotes so the value can be embedded into a PromQL label match."""
        return value.replace('"', '\\"')

    def _extract_metric_values(self, results: list) -> dict:
        """Map a PromQL ``vector`` result list to ``{metric_name: Decimal}``."""
        return {
            result["metric"]["__name__"]: Decimal(result["value"][1])
            for result in results
            if "metric" in result
            and "__name__" in result["metric"]
            and "value" in result
            and len(result["value"]) > 1
        }