"""License verification for aicert-pro using Ed25519 signatures."""

import base64
import os
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Optional

from cryptography.hazmat.primitives.asymmetric import ed25519
from cryptography.hazmat.primitives import serialization
from cryptography.exceptions import InvalidSignature

# Embedded Ed25519 public key for license verification
LICENSE_PUBLIC_KEY = ed25519.Ed25519PublicKey.from_public_bytes(
    base64.b64decode(
        "PgOgheIjMd7RMHXb4vZ48sWjp2YY0aQGe2de3H2XJc4="
    )
)

LICENSE_ENV_VAR = "AICERT_PRO_LICENSE"


class LicenseError(Exception):
    """Base exception for license errors."""

    exit_code = 5


class LicenseMissingError(LicenseError):
    """Raised when license is not found."""

    def __init__(self) -> None:
        super().__init__("License not found. Set AICERT_PRO_LICENSE environment variable.")


class LicenseInvalidError(LicenseError):
    """Raised when license signature is invalid."""

    def __init__(self) -> None:
        super().__init__("License signature is invalid.")


class LicenseExpiredError(LicenseError):
    """Raised when license has expired."""

    def __init__(self, expired_at: str) -> None:
        super().__init__(f"License expired at {expired_at}.")


@dataclass
class LicenseData:
    """License data parsed from base64-encoded JSON."""

    license_id: str
    owner: str
    issued_at: str
    expires_at: Optional[str] = None
    features: Optional[list[str]] = None
    signature: Optional[str] = None

    @property
    def is_expired(self) -> bool:
        """Check if license is expired."""
        if self.expires_at is None:
            return False
        try:
            # Handle Unix timestamp (integer) or ISO format
            if self.expires_at.isdigit():
                from datetime import datetime
                expires = datetime.fromtimestamp(int(self.expires_at), tz=timezone.utc)
                return datetime.now(timezone.utc) > expires
            else:
                expires = datetime.fromisoformat(self.expires_at.replace("Z", "+00:00"))
                return datetime.now(timezone.utc) > expires
        except ValueError:
            return False


def verify_license(license_data: LicenseData, skip_signature: bool = False) -> bool:
    """Verify Ed25519 signature on license data.
    
    Args:
        license_data: Parsed license data including signature.
        skip_signature: If True, skip signature verification (for testing).
        
    Returns:
        True if signature is valid.
        
    Raises:
        LicenseInvalidError: If signature verification fails.
        LicenseExpiredError: If license has expired.
    """
    if license_data.is_expired:
        raise LicenseExpiredError(license_data.expires_at)
    
    if skip_signature:
        return True
    
    if license_data.signature is None:
        raise LicenseInvalidError()
    
    # Reconstruct the same payload JSON that was signed (must match issue_license.py)
    payload_dict = {
        "sub": license_data.license_id,
        "plan": license_data.features[0] if license_data.features else "",
        "issued_at": int(license_data.issued_at),
        "expires_at": int(license_data.expires_at) if license_data.expires_at else None,
        "version": 1,
    }
    # Remove None values
    payload_dict = {k: v for k, v in payload_dict.items() if v is not None}
    
    # Create canonical JSON bytes (same as canonical_json_bytes in issue_license.py)
    import json
    payload_bytes = json.dumps(payload_dict, sort_keys=True, separators=(",", ":")).encode("utf-8")
    
    try:
        # Handle base64url encoding (convert -_ to +/ and add padding)
        signature = license_data.signature
        signature = signature.replace("-", "+").replace("_", "/")
        padding_needed = 4 - (len(signature) % 4)
        if padding_needed != 4:
            signature += "=" * padding_needed
        signature_bytes = base64.b64decode(signature)
        
        LICENSE_PUBLIC_KEY.verify(signature_bytes, payload_bytes)
    except (InvalidSignature, ValueError) as e:
        raise LicenseInvalidError()
    
    return True


def read_license_from_env() -> str:
    """Read license string from environment variable.
    
    Returns:
        Base64-encoded license string.
        
    Raises:
        LicenseMissingError: If environment variable is not set.
    """
    license_str = os.environ.get(LICENSE_ENV_VAR)
    if not license_str:
        raise LicenseMissingError()
    return license_str


def require_pro() -> LicenseData:
    """Require and verify a valid pro license.
    
    Returns:
        Parsed and verified license data.
        
    Exits:
        Exit code 5 if license is missing or invalid.
    """
    try:
        skip_signature = os.environ.get("AICERT_PRO_SKIP_SIGNATURE", "").lower() in ("1", "true", "yes")
        
        license_str = read_license_from_env()
        
        # Strip license prefix if present (e.g., "AICERT-1.")
        if "." in license_str:
            parts = license_str.split(".", 1)[1]
        else:
            parts = license_str
        
        # Split into JSON payload and signature
        if "." in parts:
            json_b64, signature_b64 = parts.split(".", 1)
        else:
            json_b64 = parts
            signature_b64 = None
        
        # Handle base64url encoding in JSON (convert -_ to +/)
        json_b64 = json_b64.replace("-", "+").replace("_", "/")
        
        # Add base64 padding if necessary
        padding_needed = 4 - (len(json_b64) % 4)
        if padding_needed != 4:
            json_b64 += "=" * padding_needed
        
        import json
        data = json.loads(base64.b64decode(json_b64))
        
        # Map JSON fields to LicenseData (handle different field names)
        license_data = LicenseData(
            license_id=data.get("sub", data.get("license_id", "")),
            owner=data.get("sub", data.get("owner", "")),
            issued_at=str(data.get("issued_at", "")),
            expires_at=str(data.get("expires_at", "")) if data.get("expires_at") else None,
            features=[data.get("plan", "")],
            signature=signature_b64,
        )
        verify_license(license_data, skip_signature=skip_signature)
        return license_data
    except LicenseError:
        raise
    except Exception:
        raise LicenseInvalidError()
