AegisSight-Monitor/src/services/license_service.py

"""Lizenz-Verwaltung und -Pruefung."""
import logging
import os
from datetime import datetime
from config import TIMEZONE
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")


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)
    unlimited_budget = bool(lic_dict.get("unlimited_budget"))
    credits_total = lic_dict.get("credits_total")
    credits_used = lic_dict.get("credits_used") or 0

    # 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 aufgebrauchten Credits, ausser unlimited)
    budget_exceeded = False
    if not unlimited_budget and credits_total and credits_total > 0:
        if credits_used >= credits_total:
            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": "Token-Budget aufgebraucht",
            "unlimited_budget": False,
            "credits_total": credits_total,
            "credits_used": credits_used,
        }

    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,
    }


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,
) -> None:
    """Verbucht Token-Verbrauch auf einen Tenant.

    Aktualisiert `token_usage_monthly` (UPSERT pro organization_id+year_month+source)
    und zieht Credits von der aktiven Lizenz ab (wenn cost_per_credit gesetzt).

    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' | 'enhance' | 'chat'

    Der Helper ruft KEIN db.commit() auf — die Transaktionsgrenzen bestimmt der Caller.
    Ohne Verbrauch (total_cost_usd == 0) oder ohne tenant_id wird nichts gebucht.
    """
    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

    if total_cost <= 0:
        return

    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,
        ),
    )

    lic_cursor = await db.execute(
        "SELECT cost_per_credit FROM licenses WHERE organization_id = ? AND status = 'active' ORDER BY id DESC LIMIT 1",
        (tenant_id,),
    )
    lic = await lic_cursor.fetchone()
    credits_consumed = 0.0
    if lic and lic["cost_per_credit"] and lic["cost_per_credit"] > 0:
        credits_consumed = total_cost / lic["cost_per_credit"]
        await db.execute(
            "UPDATE licenses SET credits_used = COALESCE(credits_used, 0) + ? WHERE organization_id = ? AND status = 'active'",
            (round(credits_consumed, 2), tenant_id),
        )

    logger.info(
        f"charge_usage_to_tenant[{source}] Tenant {tenant_id}: "
        f"${total_cost:.4f} -> {round(credits_consumed, 2)} Credits"
    )


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