"""License enforcement for StackSage.

Goal: allow distributing StackSage with an offline-verifiable license that expires.

License format (env/file):

    STACKSAGE1.<payload_b64url>.<sig_b64url>

Where:
- payload is JSON bytes (UTF-8)
- signature is Ed25519(signature over b"STACKSAGE1." + payload_b64url_bytes)

A public key must be available in the runtime environment (baked into the image
or provided via env var). The private signing key must be kept by StackSage.
"""

from __future__ import annotations

import base64
import json
import os
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from typing import Any


class LicenseError(RuntimeError):
    pass


@dataclass(frozen=True)
class LicenseClaims:
    customer: str | None
    plan: str | None
    expires_at: datetime
    issued_at: datetime | None


_LICENSE_ENV_VAR = "STACKSAGE_LICENSE"
_LICENSE_FILE_ENV_VAR = "STACKSAGE_LICENSE_FILE"
_PUBLIC_KEY_B64_ENV_VAR = "STACKSAGE_LICENSE_PUBLIC_KEY_B64"

# Public key bundled with the package — customers never need to set the env var.
# (This is a public key; keeping it here is safe.)
_BUNDLED_PUBLIC_KEY_B64 = "J1ayeD62DtsU6EdSFedNYJ6xHyojA9WOSTAeC45tEu4"

# Small skew window to avoid flaky failures at the boundary.
_DEFAULT_CLOCK_SKEW_SECONDS = 60

# Maximum payload size to prevent DoS (16KB should be plenty for license claims)
_MAX_PAYLOAD_SIZE_BYTES = 16384

# Maximum token length (base64 overhead + payload + signature)
_MAX_TOKEN_LENGTH = 32768


def _b64url_decode(data: str) -> bytes:
    if not data:
        return b""
    padding = "=" * (-len(data) % 4)
    return base64.urlsafe_b64decode(data + padding)


def _b64url_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")


def _load_public_key_bytes() -> bytes:
    b64 = os.environ.get(_PUBLIC_KEY_B64_ENV_VAR, "").strip()
    if not b64:
        # Fall back to key bundled with the package.
        b64 = _BUNDLED_PUBLIC_KEY_B64
    key = _b64url_decode(b64)
    if len(key) != 32:
        raise LicenseError(
            f"Invalid license public key length: expected 32 bytes, got {len(key)}."
        )
    return key


def _load_license_string() -> str:
    token = os.environ.get(_LICENSE_ENV_VAR, "").strip()
    if token:
        return token

    path = os.environ.get(_LICENSE_FILE_ENV_VAR, "").strip()
    if not path:
        raise LicenseError(
            f"Missing license. Provide env var {_LICENSE_ENV_VAR} or {_LICENSE_FILE_ENV_VAR}."
        )
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read().strip()
    except FileNotFoundError as e:
        raise LicenseError(f"License file not found: {path}") from e
    except OSError as e:
        raise LicenseError(f"Failed to read license file: {path}") from e


def _parse_payload(payload_bytes: bytes) -> dict[str, Any]:
    try:
        payload = json.loads(payload_bytes.decode("utf-8"))
    except Exception as e:
        raise LicenseError("Invalid license payload JSON") from e
    if not isinstance(payload, dict):
        raise LicenseError("Invalid license payload: expected JSON object")
    return payload


def _parse_unix_ts(value: Any, field: str) -> datetime:
    if not isinstance(value, int):
        raise LicenseError(
            f"Invalid license field {field}: expected integer unix timestamp"
        )
    try:
        return datetime.fromtimestamp(value, tz=UTC)
    except Exception as e:
        raise LicenseError(f"Invalid license field {field}") from e


def validate_license(
    *,
    now: datetime | None = None,
    clock_skew_seconds: int = _DEFAULT_CLOCK_SKEW_SECONDS,
) -> LicenseClaims:
    """Validate the configured license.

    Raises LicenseError if missing/invalid/expired.
    """

    token = _load_license_string()

    # Prevent DoS: reject excessively long tokens early
    if len(token) > _MAX_TOKEN_LENGTH:
        raise LicenseError("Invalid license: token too large")

    parts = token.split(".")
    if len(parts) != 3 or parts[0] != "STACKSAGE1":
        raise LicenseError("Invalid license format")

    payload_b64 = parts[1]
    sig_b64 = parts[2]

    # Validate base64url segments are non-empty
    if not payload_b64 or not sig_b64:
        raise LicenseError("Invalid license format: empty segments")

    payload_bytes = _b64url_decode(payload_b64)

    # Prevent DoS: reject oversized payloads
    if len(payload_bytes) > _MAX_PAYLOAD_SIZE_BYTES:
        raise LicenseError("Invalid license: payload too large")

    sig = _b64url_decode(sig_b64)
    if len(sig) != 64:
        raise LicenseError("Invalid license signature: incorrect length")

    payload = _parse_payload(payload_bytes)
    exp = _parse_unix_ts(payload.get("exp"), "exp")
    iat_val = payload.get("iat")
    iat = _parse_unix_ts(iat_val, "iat") if isinstance(iat_val, int) else None

    current = now or datetime.now(UTC)

    # Verify signature last (after cheap parsing), but before returning claims.
    try:
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
    except Exception as e:  # pragma: no cover
        raise LicenseError(
            "cryptography dependency missing; cannot verify license signature"
        ) from e

    pubkey = Ed25519PublicKey.from_public_bytes(_load_public_key_bytes())

    signing_input = ("STACKSAGE1." + payload_b64).encode("ascii")
    try:
        pubkey.verify(sig, signing_input)
    except Exception as e:
        raise LicenseError("Invalid license: signature verification failed") from e

    # Additional tamper detection: verify payload re-encodes correctly
    # This catches cases where payload_bytes was modified after signature verification
    if _b64url_encode(payload_bytes) != payload_b64:
        raise LicenseError("Invalid license: payload encoding mismatch")

    # Time validity checks (strict; small skew only for not-before).
    # Expiration is enforced strictly to avoid extending validity.
    if current >= exp:
        raise LicenseError("License expired")

    # Optional not-before support.
    nbf_val = payload.get("nbf")
    if isinstance(nbf_val, int):
        nbf = _parse_unix_ts(nbf_val, "nbf")
        if clock_skew_seconds:
            if current + timedelta(seconds=clock_skew_seconds) < nbf:
                raise LicenseError("License not yet valid")
        elif current < nbf:
            raise LicenseError("License not yet valid")

    return LicenseClaims(
        customer=(
            payload.get("customer")
            if isinstance(payload.get("customer"), str)
            else None
        ),
        plan=payload.get("plan") if isinstance(payload.get("plan"), str) else None,
        expires_at=exp,
        issued_at=iat,
    )


def enforce_license() -> LicenseClaims:
    """Convenience wrapper for callers that just want "raise on failure"."""

    return validate_license()
