# SPDX-License-Identifier: Apache-2.0
"""
License management for model encryption authorization.

Supports offline license verification via RSA signature.
This module is device-agnostic.
"""

import base64
import hashlib
import json
import os
import uuid
from dataclasses import dataclass, field
from datetime import date
from typing import TYPE_CHECKING, List, Optional

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from cryptography.hazmat.backends import default_backend
from cryptography.exceptions import InvalidSignature

if TYPE_CHECKING:
    from lightr_vllm_core.base import BaseFingerprint


@dataclass
class LicenseInfo:
    """License information structure."""
    license_id: str
    customer: str
    model_key_id: str          # hex key_id matching .model_key file
    dek_encrypted: str         # base64-encoded DEK encrypted with vendor private key
    dek_hash: str              # SHA-256 hash of plaintext DEK (for integrity)
    issued_at: str             # ISO date string
    expire_at: Optional[str]   # ISO date string, None = never expires
    fingerprint_policy: str    # FingerprintPolicy value
    allowed_fingerprints: List[str] = field(default_factory=list)
    max_nodes: int = 1
    signature: str = ""        # base64-encoded RSA signature

    def to_dict(self, include_signature: bool = True) -> dict:
        d = {
            "license_id": self.license_id,
            "customer": self.customer,
            "model_key_id": self.model_key_id,
            "dek_encrypted": self.dek_encrypted,
            "dek_hash": self.dek_hash,
            "issued_at": self.issued_at,
            "expire_at": self.expire_at,
            "fingerprint_policy": self.fingerprint_policy,
            "allowed_fingerprints": self.allowed_fingerprints,
            "max_nodes": self.max_nodes,
        }
        if include_signature:
            d["signature"] = self.signature
        return d

    @classmethod
    def from_dict(cls, d: dict) -> "LicenseInfo":
        return cls(
            license_id=d["license_id"],
            customer=d["customer"],
            model_key_id=d["model_key_id"],
            dek_encrypted=d["dek_encrypted"],
            dek_hash=d["dek_hash"],
            issued_at=d["issued_at"],
            expire_at=d.get("expire_at"),
            fingerprint_policy=d.get("fingerprint_policy", "none"),
            allowed_fingerprints=d.get("allowed_fingerprints", []),
            max_nodes=d.get("max_nodes", 1),
            signature=d.get("signature", ""),
        )

    def signing_payload(self) -> bytes:
        """Get canonical bytes to sign (excludes signature field)."""
        d = self.to_dict(include_signature=False)
        return json.dumps(d, sort_keys=True, separators=(',', ':')).encode('utf-8')


class LicenseManager:
    """Creates and signs license files (vendor-side tool)."""

    def __init__(self, vendor_private_key_path: str):
        with open(vendor_private_key_path, 'rb') as f:
            self.private_key = serialization.load_pem_private_key(
                f.read(), password=None, backend=default_backend()
            )

    def create_license(
        self,
        customer: str,
        model_key_id: str,
        dek: bytes,
        vendor_secret: bytes,
        fingerprint_policy: str = "none",
        allowed_fingerprints: Optional[List[str]] = None,
        expire_at: Optional[str] = None,
        max_nodes: int = 1,
    ) -> LicenseInfo:
        """
        Create and sign a license.

        Args:
            customer: Customer/organization name
            model_key_id: Hex key_id from encrypted model
            dek: Plaintext DEK bytes
            vendor_secret: Secret bytes for DEK encryption
            fingerprint_policy: Machine binding policy ("none", "single", "cluster")
            allowed_fingerprints: List of allowed machine fingerprints
            expire_at: Expiry date as ISO string, None = no expiry
            max_nodes: Maximum number of nodes allowed

        Returns:
            Signed LicenseInfo
        """
        license_id = str(uuid.uuid4())
        dek_hash = hashlib.sha256(dek).hexdigest()

        dek_key = OfflineLicenseManager.derive_dek_key(vendor_secret, license_id)
        dek_cipher = OfflineLicenseManager.encrypt_dek_for_license(dek, dek_key)

        info = LicenseInfo(
            license_id=license_id,
            customer=customer,
            model_key_id=model_key_id,
            dek_encrypted=dek_cipher,
            dek_hash=dek_hash,
            issued_at=date.today().isoformat(),
            expire_at=expire_at,
            fingerprint_policy=fingerprint_policy,
            allowed_fingerprints=allowed_fingerprints or [],
            max_nodes=max_nodes,
        )

        # Sign the payload
        signature = self.private_key.sign(
            info.signing_payload(),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH,
            ),
            hashes.SHA256(),
        )
        info.signature = base64.b64encode(signature).decode()
        return info

    def save_license(self, info: LicenseInfo, path: str) -> None:
        """Save license to JSON file."""
        with open(path, 'w') as f:
            json.dump(info.to_dict(), f, indent=2)

    @staticmethod
    def generate_keypair(
        private_key_path: str,
        public_key_path: str,
        key_size: int = 3072,
        passphrase: Optional[bytes] = None,
    ) -> None:
        """Generate RSA keypair for license signing."""
        key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=key_size,
            backend=default_backend(),
        )
        if passphrase:
            encryption_algorithm = serialization.BestAvailableEncryption(passphrase)
        else:
            encryption_algorithm = serialization.NoEncryption()
        with open(private_key_path, 'wb') as f:
            f.write(key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=encryption_algorithm,
            ))
        with open(public_key_path, 'wb') as f:
            f.write(key.public_key().public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo,
            ))


class LicenseValidator:
    """Validates licenses at runtime (user-side)."""

    def __init__(self, vendor_public_key_path: str):
        with open(vendor_public_key_path, 'rb') as f:
            self.public_key = serialization.load_pem_public_key(
                f.read(), backend=default_backend()
            )

    def load_license(self, license_path: str) -> LicenseInfo:
        """Load license from file."""
        with open(license_path, 'r') as f:
            return LicenseInfo.from_dict(json.load(f))

    def validate(
        self,
        info: LicenseInfo,
        fingerprint_provider: Optional["BaseFingerprint"] = None,
    ) -> None:
        """
        Validate license. Raises on any failure.

        Args:
            info: License to validate
            fingerprint_provider: Platform-specific fingerprint provider

        Raises:
            ValueError: If license is invalid, expired, or machine doesn't match
        """
        # 1. Verify RSA signature
        self._verify_signature(info)

        # 2. Check expiry
        self._check_expiry(info)

        # 3. Check machine fingerprint (if policy requires)
        if fingerprint_provider:
            self._check_fingerprint(info, fingerprint_provider)

    def extract_dek_from_model_dir(
        self,
        info: LicenseInfo,
        model_dir: str,
        vendor_secret: Optional[bytes] = None,
    ) -> bytes:
        """
        Extract DEK from model directory using license info.

        Args:
            info: Validated license info
            model_dir: Path to model directory containing .model_key
            vendor_secret: Vendor secret for DEK decryption

        Returns:
            Plaintext DEK bytes
        """
        model_key_path = os.path.join(model_dir, ".model_key")
        if not os.path.isfile(model_key_path):
            raise ValueError(f"Model key file not found: {model_key_path}")

        with open(model_key_path, 'r') as f:
            key_data = json.load(f)

        # Verify key_id matches
        if key_data.get("key_id") != info.model_key_id:
            raise ValueError("License key_id mismatch. This license is not for this model.")

        # Extract and decrypt DEK
        if vendor_secret:
            dek_key = OfflineLicenseManager.derive_dek_key(
                vendor_secret, info.license_id
            )
            dek = OfflineLicenseManager.decrypt_dek_from_license(
                info.dek_encrypted, dek_key
            )
        else:
            raise ValueError("vendor_secret is required for DEK decryption.")

        # Verify DEK hash
        dek_hash = hashlib.sha256(dek).hexdigest()
        if dek_hash != info.dek_hash:
            raise ValueError("DEK integrity check failed. License or model may be corrupted.")

        return dek

    def _verify_signature(self, info: LicenseInfo) -> None:
        """Verify RSA signature on license payload."""
        try:
            self.public_key.verify(
                base64.b64decode(info.signature),
                info.signing_payload(),
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH,
                ),
                hashes.SHA256(),
            )
        except InvalidSignature:
            raise ValueError("License signature verification failed. License may be tampered.")

    def _check_expiry(self, info: LicenseInfo) -> None:
        """Check if license is expired."""
        if info.expire_at is None:
            return
        expire_date = date.fromisoformat(info.expire_at)
        if date.today() > expire_date:
            raise ValueError(f"License expired on {info.expire_at}.")

    def check_expiry_with_grace(self, info: LicenseInfo, grace_days: int = 3) -> str:
        """
        Check license expiry, returning status string.

        Returns:
            "valid"   — not expired
            "warning" — expired but within grace period
            "expired" — expired beyond grace period
        """
        if info.expire_at is None:
            return "valid"
        expire_date = date.fromisoformat(info.expire_at)
        today = date.today()
        if today <= expire_date:
            return "valid"
        days_over = (today - expire_date).days
        if days_over <= grace_days:
            return "warning"
        return "expired"

    def _check_fingerprint(
        self,
        info: LicenseInfo,
        fingerprint_provider: "BaseFingerprint",
    ) -> None:
        """Check machine fingerprint against license."""
        from lightr_vllm_core.base import FingerprintPolicy

        policy = FingerprintPolicy(info.fingerprint_policy)

        if policy == FingerprintPolicy.NONE:
            return

        current_fp = fingerprint_provider.get_fingerprint()

        if policy == FingerprintPolicy.SINGLE:
            if not info.allowed_fingerprints:
                raise ValueError("License has SINGLE policy but no allowed fingerprints.")
            if current_fp not in info.allowed_fingerprints:
                raise ValueError(
                    "Machine fingerprint mismatch. "
                    "This license is not authorized for this machine."
                )

        elif policy == FingerprintPolicy.CLUSTER:
            if not info.allowed_fingerprints:
                raise ValueError("License has CLUSTER policy but no allowed fingerprints.")
            if current_fp not in info.allowed_fingerprints:
                raise ValueError(
                    "Machine fingerprint mismatch. "
                    "This machine is not in the authorized cluster."
                )


class OfflineLicenseManager:
    """Simplified offline license manager for DEK encryption."""

    @staticmethod
    def derive_dek_key(vendor_secret: bytes, license_id: str) -> bytes:
        """Derive a key for encrypting the DEK in the license."""
        from cryptography.hazmat.primitives.kdf.hkdf import HKDF
        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=license_id.encode(),
            info=b"vllm-security-dek-key",
            backend=default_backend(),
        )
        return hkdf.derive(vendor_secret)

    @staticmethod
    def encrypt_dek_for_license(dek: bytes, dek_key: bytes) -> str:
        """Encrypt DEK using derived key."""
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        aesgcm = AESGCM(dek_key)
        nonce = os.urandom(12)
        encrypted = aesgcm.encrypt(nonce, dek, None)
        return base64.b64encode(nonce + encrypted).decode()

    @staticmethod
    def decrypt_dek_from_license(dek_encrypted_b64: str, dek_key: bytes) -> bytes:
        """Decrypt DEK from license."""
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        data = base64.b64decode(dek_encrypted_b64)
        nonce = data[:12]
        ciphertext = data[12:]
        aesgcm = AESGCM(dek_key)
        return aesgcm.decrypt(nonce, ciphertext, None)