AegisSight-Monitor/src/services/license_service.py

"""Lizenz-Verwaltung und -Pruefung."""
import logging
import os
from datetime import datetime
from config import TIMEZONE, BILLING_MODE, CREDIT_TARIFF
import aiosqlite

logger = logging.getLogger("osint.license")


def _staging_mode() -> bool:
    """Staging-Mode aktiv? Wenn ja, gilt: immer unlimited Budget, kein Hard-Stop.

    Wird ueber ENV-Variable STAGING_MODE=1 (oder true) aktiviert.
    Nur in Staging-.env gesetzt; Live-.env hat das Flag nicht.
    """
    return os.environ.get("STAGING_MODE", "").lower() in ("1", "true", "yes")


def _current_period() -> str:
    """Kennung der laufenden Abrechnungsperiode (Kalendermonat)."""
    return datetime.now(TIMEZONE).strftime("%Y-%m")


def _tariff_key(source: str, incident_type: str | None) -> str:
    """Abrechnungsquelle auf einen Tarifschluessel abbilden.

    Ein Refresh kostet je nach Lagentyp unterschiedlich viel, deshalb wird
    'monitor' anhand des Typs aufgeteilt. Alle anderen Quellen entsprechen
    direkt einem Schluessel in CREDIT_TARIFF.
    """
    if source == "monitor":
        return "monitor_research" if incident_type == "research" else "monitor_adhoc"
    return source


async def roll_credit_period(db: aiosqlite.Connection, lic: dict) -> dict:
    """Setzt das Guthaben zurueck, wenn eine neue Abrechnungsperiode begonnen hat.

    Wird traege bei jeder Lizenzpruefung aufgerufen statt ueber einen Zeitplan.
    Das ist robuster, weil ein verpasster Monatswechsel beim naechsten Zugriff
    ohnehin nachgeholt wird und ohne Nutzung auch nichts verbraucht wird.

    Bei aktivem Uebertrag wandert ungenutztes Guthaben in die Folgeperiode,
    gedeckelt auf ein Monatskontingent, damit es nicht unbegrenzt anwaechst.

    Returns:
        Das ggf. aktualisierte Lizenz-dict (in-place ergaenzt).
    """
    if (lic.get("credits_period") or "monthly") != "monthly":
        return lic
    if not lic.get("credits_total"):
        return lic

    period = _current_period()
    started = lic.get("credits_period_start")

    if not started:
        # Bestandslizenz ohne Periodenmarke. Marke setzen, Verbrauch stehen
        # lassen -- ein Reset wuerde dem Kunden hier Guthaben schenken, das er
        # in diesem Monat schon verbraucht hat.
        await db.execute(
            "UPDATE licenses SET credits_period_start = ? WHERE id = ?",
            (period, lic["id"]),
        )
        await db.commit()
        lic["credits_period_start"] = period
        return lic

    if started == period:
        return lic

    total = lic.get("credits_total") or 0
    carried_old = lic.get("credits_carried") or 0
    used = lic.get("credits_used") or 0

    if lic.get("credits_rollover"):
        carried_new = max(0.0, (total + carried_old) - used)
        carried_new = min(carried_new, float(total))  # hoechstens ein Monat
    else:
        carried_new = 0.0

    await db.execute(
        """UPDATE licenses
           SET credits_used = 0, credits_carried = ?, credits_period_start = ?,
               budget_warning_sent = 0
           WHERE id = ?""",
        (round(carried_new, 2), period, lic["id"]),
    )
    await db.commit()

    lic["credits_used"] = 0
    lic["credits_carried"] = carried_new
    lic["credits_period_start"] = period
    lic["budget_warning_sent"] = 0

    logger.info(
        f"Lizenz {lic['id']}: neue Periode {period}, Verbrauch zurueckgesetzt "
        f"(Uebertrag {round(carried_new, 2)} Einheiten)"
    )
    return lic


async def _notify_budget_warning(
    db: aiosqlite.Connection, organization_id: int, percent: float, remaining: float
) -> None:
    """Legt eine Warnung an, sobald die Schwelle des Guthabens erreicht ist.

    Die Meldung geht an alle aktiven Nutzer der Organisation, die sich schon
    einmal angemeldet haben. Ein E-Mail-Versand haengt hier bewusst nicht dran,
    das waere ein eigener Schritt ueber email_utils.
    """
    cursor = await db.execute(
        "SELECT id FROM users WHERE organization_id = ? AND is_active = 1 AND last_login_at IS NOT NULL",
        (organization_id,),
    )
    user_ids = [row["id"] for row in await cursor.fetchall()]
    now = datetime.now(TIMEZONE).strftime("%Y-%m-%d %H:%M:%S")

    title = "Credits zu {p} Prozent verbraucht".format(p=int(percent))
    text = (
        f"Es sind noch rund {int(remaining)} Credits in dieser Abrechnungsperiode "
        f"verfuegbar. Sind die Credits aufgebraucht, bleiben die Lagen lesbar, es "
        f"lassen sich aber keine neuen Aktualisierungen mehr starten."
    )

    for user_id in user_ids:
        await db.execute(
            """INSERT INTO notifications (user_id, incident_id, type, title, text, icon, tenant_id, created_at)
               VALUES (?, NULL, 'budget_warning', ?, ?, 'warning', ?, ?)""",
            (user_id, title, text, organization_id, now),
        )

    logger.info(
        f"Budget-Warnung fuer Org {organization_id} an {len(user_ids)} Nutzer "
        f"({int(percent)} Prozent verbraucht)"
    )


async def check_license(db: aiosqlite.Connection, organization_id: int) -> dict:
    """Prueft den Lizenzstatus einer Organisation.

    Returns:
        dict mit: valid, status, license_type, max_users, current_users, read_only,
        read_only_reason, message, unlimited_budget, credits_total, credits_used
    """
    # Organisation pruefen
    cursor = await db.execute(
        "SELECT id, name, is_active FROM organizations WHERE id = ?",
        (organization_id,),
    )
    org = await cursor.fetchone()
    if not org:
        return {"valid": False, "status": "not_found", "read_only": True,
                "read_only_reason": "not_found",
                "message": "Organisation nicht gefunden"}

    if not org["is_active"]:
        return {"valid": False, "status": "org_disabled", "read_only": True,
                "read_only_reason": "org_disabled",
                "message": "Organisation deaktiviert"}

    # Aktive Lizenz suchen
    cursor = await db.execute(
        """SELECT * FROM licenses
           WHERE organization_id = ? AND status = 'active'
           ORDER BY created_at DESC LIMIT 1""",
        (organization_id,),
    )
    license_row = await cursor.fetchone()

    if not license_row:
        return {"valid": False, "status": "no_license", "read_only": True,
                "read_only_reason": "no_license",
                "message": "Keine aktive Lizenz"}

    # Felder zur weiteren Verwendung extrahieren
    lic_dict = dict(license_row)

    # Periodenwechsel nachholen, bevor irgendetwas geprueft wird. Sonst haengt
    # ein Kunde mit monatlichem Kontingent im Nur-Lese-Modus fest, obwohl der
    # neue Monat laengst begonnen hat.
    lic_dict = await roll_credit_period(db, lic_dict)

    unlimited_budget = bool(lic_dict.get("unlimited_budget"))
    credits_total = lic_dict.get("credits_total")
    credits_used = lic_dict.get("credits_used") or 0
    credits_carried = lic_dict.get("credits_carried") or 0
    credits_period = lic_dict.get("credits_period") or "monthly"
    # Verfuegbar ist das Kontingent plus ein etwaiger Uebertrag aus dem Vormonat.
    credits_available = (credits_total or 0) + credits_carried

    # STAGING_MODE: kein Token-Budget-Hard-Stop, immer unlimited
    if _staging_mode():
        unlimited_budget = True

    # Ablauf pruefen
    now = datetime.now(TIMEZONE)
    valid_until = license_row["valid_until"]

    if valid_until is not None:
        try:
            expiry = datetime.fromisoformat(valid_until)
            if expiry.tzinfo is None:
                expiry = expiry.replace(tzinfo=TIMEZONE)
            if now > expiry:
                return {
                    "valid": False,
                    "status": "expired",
                    "license_type": license_row["license_type"],
                    "read_only": True,
                    "read_only_reason": "expired",
                    "message": "Lizenz abgelaufen",
                    "unlimited_budget": unlimited_budget,
                    "credits_total": credits_total,
                    "credits_used": credits_used,
                }
        except (ValueError, TypeError):
            pass

    # Budget-Check (Hard-Stop bei aufgebrauchtem Guthaben, ausser unlimited).
    # Geprueft wird gegen das verfuegbare Guthaben, also inklusive Uebertrag.
    budget_exceeded = False
    if not unlimited_budget and credits_available > 0:
        if credits_used >= credits_available:
            budget_exceeded = True

    # Nutzerzahl pruefen
    cursor = await db.execute(
        "SELECT COUNT(*) as cnt FROM users WHERE organization_id = ? AND is_active = 1",
        (organization_id,),
    )
    current_users = (await cursor.fetchone())["cnt"]

    if budget_exceeded:
        return {
            "valid": True,  # Lizenz ist gueltig, aber Budget aufgebraucht -> read-only
            "status": "budget_exceeded",
            "license_type": license_row["license_type"],
            "max_users": license_row["max_users"],
            "current_users": current_users,
            "read_only": True,
            "read_only_reason": "budget_exceeded",
            "message": "Credits aufgebraucht",
            "unlimited_budget": False,
            "credits_total": credits_total,
            "credits_used": credits_used,
            "credits_carried": credits_carried,
            "credits_available": credits_available,
            "credits_period": credits_period,
        }

    return {
        "valid": True,
        "status": license_row["status"],
        "license_type": license_row["license_type"],
        "max_users": license_row["max_users"],
        "current_users": current_users,
        "read_only": False,
        "read_only_reason": None,
        "message": "Lizenz aktiv",
        "unlimited_budget": unlimited_budget,
        "credits_total": credits_total,
        "credits_used": credits_used,
        "credits_carried": credits_carried,
        "credits_available": credits_available,
        "credits_period": credits_period,
    }


async def can_add_user(db: aiosqlite.Connection, organization_id: int) -> tuple[bool, str]:
    """Prueft ob ein neuer Nutzer hinzugefuegt werden kann (Nutzer-Limit).

    Returns:
        (erlaubt, grund)
    """
    lic = await check_license(db, organization_id)
    if not lic["valid"]:
        return False, lic["message"]

    if lic["current_users"] >= lic["max_users"]:
        return False, f"Nutzer-Limit erreicht ({lic['current_users']}/{lic['max_users']})"

    return True, ""


async def charge_usage_to_tenant(
    db: aiosqlite.Connection,
    tenant_id: int | None,
    usage,
    source: str,
    incident_type: str | None = None,
) -> None:
    """Verbucht eine Aktion auf einen Tenant.

    Zwei getrennte Vorgaenge. `token_usage_monthly` bekommt immer die echten
    Tokenmengen und Kosten, das ist die interne Kostenkontrolle. Das Guthaben
    der Lizenz wird je nach BILLING_MODE belastet.

    'flat'   zieht den festen Satz aus CREDIT_TARIFF ab. Der Kunde kann seinen
             Verbrauch damit vorher ausrechnen, und eine spaetere Verbilligung
             des Modell-Backends veraendert sein Kontingent nicht.
    'actual' zieht die echten Kosten geteilt durch cost_per_credit ab, also das
             bisherige Verhalten.

    Args:
        db: offene aiosqlite.Connection
        tenant_id: Organisations-ID oder None (dann nur geloggt, keine DB-Buchung)
        usage: ClaudeUsage oder UsageAccumulator mit input_tokens/output_tokens/
               cache_creation_tokens/cache_read_tokens/total_cost_usd/call_count
        source: 'monitor' | 'analysis' | 'factcheck' | 'chat' | 'enhance' | 'globe'
        incident_type: 'adhoc' | 'research', nur bei source='monitor' relevant

    Der Helper ruft KEIN db.commit() auf — die Transaktionsgrenzen bestimmt der Caller.
    Ausnahme ist der Periodenwechsel, der eine eigene Transaktion braucht.
    """
    total_cost = getattr(usage, "total_cost_usd", None)
    if total_cost is None:
        total_cost = getattr(usage, "cost_usd", 0.0)

    if not tenant_id:
        logger.info(
            f"charge_usage_to_tenant[{source}]: kein tenant_id, uebersprungen "
            f"(cost=${total_cost:.4f})"
        )
        return

    # Ohne echte Kosten gibt es nichts zu statistisch erfassen. Die Guthaben-
    # Buchung laeuft im Pauschalmodus trotzdem, weil der Kunde die Aktion
    # bezahlt und nicht unseren Einkauf. Auf lokalen Modellen ist total_cost 0.
    if total_cost > 0:
        await _record_usage_statistics(db, tenant_id, usage, source, total_cost)

    await _charge_credits(db, tenant_id, source, incident_type, total_cost)


async def _record_usage_statistics(
    db: aiosqlite.Connection, tenant_id: int, usage, source: str, total_cost: float
) -> None:
    """Schreibt die echten Tokenmengen und Kosten nach token_usage_monthly."""
    input_tokens = getattr(usage, "input_tokens", 0)
    output_tokens = getattr(usage, "output_tokens", 0)
    cache_creation = getattr(usage, "cache_creation_tokens", 0)
    cache_read = getattr(usage, "cache_read_tokens", 0)
    api_calls = getattr(usage, "call_count", 1)
    refresh_increment = 1 if source == "monitor" else 0

    year_month = datetime.now(TIMEZONE).strftime("%Y-%m")

    await db.execute(
        """
        INSERT INTO token_usage_monthly
            (organization_id, year_month, source, input_tokens, output_tokens,
             cache_creation_tokens, cache_read_tokens, total_cost_usd, api_calls, refresh_count)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(organization_id, year_month, source) DO UPDATE SET
            input_tokens = input_tokens + excluded.input_tokens,
            output_tokens = output_tokens + excluded.output_tokens,
            cache_creation_tokens = cache_creation_tokens + excluded.cache_creation_tokens,
            cache_read_tokens = cache_read_tokens + excluded.cache_read_tokens,
            total_cost_usd = total_cost_usd + excluded.total_cost_usd,
            api_calls = api_calls + excluded.api_calls,
            refresh_count = refresh_count + excluded.refresh_count,
            updated_at = CURRENT_TIMESTAMP
        """,
        (
            tenant_id, year_month, source,
            input_tokens, output_tokens, cache_creation, cache_read,
            round(total_cost, 7), api_calls, refresh_increment,
        ),
    )


async def _charge_credits(
    db: aiosqlite.Connection,
    tenant_id: int,
    source: str,
    incident_type: str | None,
    total_cost: float,
) -> None:
    """Belastet das Guthaben der aktiven Lizenz und prueft die Warnschwelle."""
    lic_cursor = await db.execute(
        "SELECT * FROM licenses WHERE organization_id = ? AND status = 'active' ORDER BY id DESC LIMIT 1",
        (tenant_id,),
    )
    lic_row = await lic_cursor.fetchone()
    if not lic_row:
        return

    lic = dict(lic_row)
    if not lic.get("credits_total"):
        return  # Lizenz ohne Kontingent, nichts zu belasten

    # Periodenwechsel nachholen, bevor gebucht wird. Sonst landet der erste
    # Verbrauch des neuen Monats noch auf dem alten Zaehler.
    lic = await roll_credit_period(db, lic)

    key = _tariff_key(source, incident_type)

    if BILLING_MODE == "flat":
        credits_consumed = CREDIT_TARIFF.get(key)
        if credits_consumed is None:
            # Unbekannte Quelle. Lieber auf die echte Rechnung zurueckfallen als
            # stillschweigend gratis abzugeben.
            logger.warning(
                f"Kein Tarif fuer '{key}', falle auf tatsaechliche Kosten zurueck"
            )
            credits_consumed = _actual_credits(lic, total_cost)
    else:
        credits_consumed = _actual_credits(lic, total_cost)

    if credits_consumed <= 0:
        return

    await db.execute(
        "UPDATE licenses SET credits_used = COALESCE(credits_used, 0) + ? WHERE id = ?",
        (round(credits_consumed, 2), lic["id"]),
    )

    used_new = (lic.get("credits_used") or 0) + credits_consumed
    available = (lic.get("credits_total") or 0) + (lic.get("credits_carried") or 0)

    logger.info(
        f"charge_usage_to_tenant[{key}] Tenant {tenant_id}: "
        f"${total_cost:.4f} -> {round(credits_consumed, 2)} Einheiten "
        f"({round(used_new, 1)}/{round(available, 1)})"
    )

    await _check_budget_warning(db, tenant_id, lic, used_new, available)


def _actual_credits(lic: dict, total_cost: float) -> float:
    """Echte Kosten in Einheiten umrechnen (Modus 'actual' und Rueckfallebene)."""
    cost_per_credit = lic.get("cost_per_credit")
    if not cost_per_credit or cost_per_credit <= 0:
        return 0.0
    return total_cost / cost_per_credit


async def _check_budget_warning(
    db: aiosqlite.Connection, tenant_id: int, lic: dict, used: float, available: float
) -> None:
    """Meldet einmal je Periode, wenn die Warnschwelle erreicht ist.

    Die Schwelle steht als budget_warning_percent auf der Lizenz und war bisher
    zwar als Spalte vorhanden, wurde aber nirgends ausgewertet. Ohne sie liefen
    Kunden ohne Vorwarnung in den Nur-Lese-Modus.
    """
    if available <= 0 or lic.get("budget_warning_sent"):
        return

    threshold = lic.get("budget_warning_percent") or 80
    percent = (used / available) * 100
    if percent < threshold:
        return

    await db.execute(
        "UPDATE licenses SET budget_warning_sent = 1 WHERE id = ?", (lic["id"],)
    )
    try:
        await _notify_budget_warning(db, tenant_id, percent, max(0.0, available - used))
    except Exception as e:
        # Eine fehlgeschlagene Benachrichtigung darf die Buchung nicht kippen.
        logger.warning(f"Budget-Warnung konnte nicht zugestellt werden: {e}")


async def expire_licenses(db: aiosqlite.Connection):
    """Setzt abgelaufene Lizenzen auf 'expired'. Taeglich aufrufen."""
    cursor = await db.execute(
        """SELECT id, organization_id FROM licenses
           WHERE status = 'active'
           AND valid_until IS NOT NULL
           AND valid_until < datetime('now')"""
    )
    expired = await cursor.fetchall()

    count = 0
    for lic in expired:
        await db.execute(
            "UPDATE licenses SET status = 'expired' WHERE id = ?",
            (lic["id"],),
        )
        count += 1
        logger.info(f"Lizenz {lic['id']} fuer Org {lic['organization_id']} als abgelaufen markiert")

    if count > 0:
        await db.commit()
        logger.info(f"{count} Lizenz(en) als abgelaufen markiert")

    return count