"""
Module signing functionality for HLA-Compass SDK.

Provides RSA-PSS signing and verification for module packages.
Compatible with frontend Web Crypto API implementation.
"""

import base64
import binascii
import hashlib
import json
import logging
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, cast

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

from .config import Config

logger = logging.getLogger(__name__)


_SIGNATURE_KEYS = frozenset(
    {
        "signature",
        "publicKey",
        "public_key",
        "signatureAlgorithm",
        "signature_algorithm",
        "hashAlgorithm",
        "hash_algorithm",
        "keyFingerprint",
        "key_fingerprint",
    }
)


def _strip_signature_keys(value: Any) -> Any:
    """Recursively strip signature metadata keys from a manifest structure."""
    if isinstance(value, dict):
        return {k: _strip_signature_keys(v) for k, v in value.items() if k not in _SIGNATURE_KEYS}
    if isinstance(value, list):
        return [_strip_signature_keys(item) for item in value]
    return value


class ModuleSigner:
    """Handles RSA-PSS signing and verification of module manifests."""

    KEY_SIZE = 4096
    ALGORITHM = "RSA-PSS"
    HASH_ALGORITHM = "SHA-256"

    _SIGNATURE_KEYS = _SIGNATURE_KEYS

    def __init__(self, keys_dir: Optional[Path] = None):
        """
        Initialize the ModuleSigner.

        Args:
            keys_dir: Directory to store keys. Defaults to ``<config-dir>/keys/``.
        """
        if keys_dir is None:
            keys_dir = Config.get_config_dir() / "keys"
        self.keys_dir = Path(keys_dir)
        self.keys_dir.mkdir(parents=True, exist_ok=True)

        self.private_key_path = self.keys_dir / "private.pem"
        self.public_key_path = self.keys_dir / "public.pem"

    def generate_keys(self, force: bool = False) -> Tuple[str, str]:
        """
        Generate RSA key pair for module signing.

        Args:
            force: If True, overwrite existing keys

        Returns:
            Tuple of (private_key_path, public_key_path)
        """
        if self.private_key_path.exists() and not force:
            raise FileExistsError(f"Keys already exist at {self.keys_dir}. Use force=True to regenerate.")

        # Generate private key
        private_key = rsa.generate_private_key(public_exponent=65537, key_size=self.KEY_SIZE, backend=default_backend())

        # Save private key
        private_pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        self.private_key_path.write_bytes(private_pem)
        self.private_key_path.chmod(0o600)  # Secure permissions

        # Extract and save public key
        public_key = private_key.public_key()
        public_pem = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        self.public_key_path.write_bytes(public_pem)
        self.public_key_path.chmod(0o644)

        return str(self.private_key_path), str(self.public_key_path)

    def load_private_key(self) -> rsa.RSAPrivateKey:
        """Load private key from file."""
        if not self.private_key_path.exists():
            raise FileNotFoundError(
                f"Private key not found at {self.private_key_path}. Run 'hla-compass keys init' to generate keys."
            )

        private_pem = self.private_key_path.read_bytes()
        return cast(
            rsa.RSAPrivateKey, serialization.load_pem_private_key(private_pem, password=None, backend=default_backend())
        )

    def load_public_key(self) -> rsa.RSAPublicKey:
        """Load public key from file."""
        if not self.public_key_path.exists():
            raise FileNotFoundError(
                f"Public key not found at {self.public_key_path}. Run 'hla-compass keys init' to generate keys."
            )

        public_pem = self.public_key_path.read_bytes()
        return cast(rsa.RSAPublicKey, serialization.load_pem_public_key(public_pem, backend=default_backend()))

    def get_public_key_string(self) -> str:
        """Get public key as base64 string (without PEM headers)."""
        public_key = self.load_public_key()
        public_bytes = public_key.public_bytes(
            encoding=serialization.Encoding.DER,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        return base64.b64encode(public_bytes).decode("utf-8")

    def sign_manifest(self, manifest: Dict[str, Any]) -> str:
        """
        Sign a module manifest using RSA-PSS with SHA-256.

        Args:
            manifest: Module manifest dictionary

        Returns:
            Base64-encoded signature string
        """

        manifest_copy = _strip_signature_keys(dict(manifest))

        # Sort keys for consistent signing (must match backend canonicalization)
        manifest_json = json.dumps(manifest_copy, sort_keys=True, separators=(",", ":"))
        manifest_bytes = manifest_json.encode("utf-8")

        # Load private key
        private_key = self.load_private_key()

        # Sign using RSA-PSS with SHA-256
        signature = private_key.sign(
            manifest_bytes,
            padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
            hashes.SHA256(),
        )

        # Return base64-encoded signature
        return base64.b64encode(signature).decode("utf-8")

    def verify_signature(
        self,
        manifest: Dict[str, Any],
        signature: str,
        public_key_pem: Optional[str] = None,
    ) -> bool:
        """
        Verify a module manifest signature.

        Args:
            manifest: Module manifest dictionary
            signature: Base64-encoded signature string
            public_key_pem: Optional PEM-encoded public key (uses stored key if not provided)

        Returns:
            True if signature is valid, False otherwise
        """
        try:
            if public_key_pem:
                public_key = cast(
                    rsa.RSAPublicKey,
                    serialization.load_pem_public_key(public_key_pem.encode("utf-8"), backend=default_backend()),
                )
            else:
                public_key = self.load_public_key()

            manifest_copy = _strip_signature_keys(dict(manifest))
            manifest_json = json.dumps(manifest_copy, sort_keys=True, separators=(",", ":"))
            manifest_bytes = manifest_json.encode("utf-8")
            signature_bytes = base64.b64decode(signature)

            public_key.verify(
                signature_bytes,
                manifest_bytes,
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH,
                ),
                hashes.SHA256(),
            )
            return True
        except (InvalidSignature, ValueError, TypeError, binascii.Error):
            return False

    def get_key_fingerprint(self) -> str:
        """Get SHA-256 fingerprint of the public key."""
        public_key = self.load_public_key()
        public_bytes = public_key.public_bytes(
            encoding=serialization.Encoding.DER,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        fingerprint = hashlib.sha256(public_bytes).digest()
        return base64.b64encode(fingerprint).decode("utf-8")

    def export_public_key(self, output_path: Optional[Path] = None) -> str:
        """
        Export public key in PEM format.

        Args:
            output_path: Optional path to save the key

        Returns:
            PEM-encoded public key string
        """
        public_pem = self.public_key_path.read_text()

        if output_path:
            output_path = Path(output_path)
            output_path.write_text(public_pem)

        return public_pem


def sign_module_package(package_path: Path, signer: Optional[ModuleSigner] = None) -> Dict[str, Any]:
    """
    Sign a module package by adding signature to manifest.

    Args:
        package_path: Path to module package directory
        signer: Optional ModuleSigner instance (creates new if not provided)

    Returns:
        Updated manifest with signature
    """
    if signer is None:
        signer = ModuleSigner()

    # Load manifest
    manifest_path = package_path / "manifest.json"
    if not manifest_path.exists():
        raise FileNotFoundError(f"Manifest not found at {manifest_path}")

    with open(manifest_path, "r", encoding="utf-8") as f:
        manifest = json.load(f)

    # Sign manifest
    signature = signer.sign_manifest(manifest)

    # Add signature and public key to manifest
    manifest["signature"] = signature
    manifest["publicKey"] = signer.get_public_key_string()
    manifest["signatureAlgorithm"] = signer.ALGORITHM
    manifest["hashAlgorithm"] = signer.HASH_ALGORITHM
    manifest["keyFingerprint"] = signer.get_key_fingerprint()

    # Save updated manifest
    with open(manifest_path, "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2)

    return manifest


def verify_module_package(
    package_path: Path,
    public_key_pem: Optional[str] = None,
    *,
    allow_self_signed: bool = False,
) -> bool:
    """
    Verify a signed module package.

    Args:
        package_path: Path to module package directory
        public_key_pem: Optional PEM-encoded public key (trusted external key)
        allow_self_signed: If True, fall back to the embedded publicKey in the
            manifest when no external key is provided.  This only proves
            internal consistency (not authenticity) and should only be used
            during local development.

    Returns:
        True if signature is valid, False otherwise
    """
    # Load manifest
    manifest_path = package_path / "manifest.json"
    if not manifest_path.exists():
        return False

    with open(manifest_path, "r", encoding="utf-8") as f:
        manifest = json.load(f)

    # Check for signature
    if "signature" not in manifest:
        return False

    # Use embedded public key only when explicitly allowed (dev workflows)
    if not public_key_pem and "publicKey" in manifest:
        if not allow_self_signed:
            logger.warning(
                "Cannot verify signature: no trusted public key provided. "
                "The manifest contains an embedded publicKey but self-signed "
                "verification is disabled. Pass allow_self_signed=True for "
                "local development or provide a trusted public_key_pem."
            )
            return False
        try:
            public_der = base64.b64decode(manifest["publicKey"])
            public_key = serialization.load_der_public_key(public_der, backend=default_backend())
            public_key_pem = public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo,
            ).decode("utf-8")
        except (ValueError, TypeError, binascii.Error):
            logger.warning("Manifest contains an invalid embedded publicKey")
            return False

    signer = ModuleSigner()
    return signer.verify_signature(manifest, manifest["signature"], public_key_pem)
