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

from __future__ import annotations

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:
    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"

    def __init__(self, keys_dir: Optional[Path] = None):
        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]:
        if self.private_key_path.exists() and not force:
            raise FileExistsError(f"Keys already exist at {self.keys_dir}. Use force=True to regenerate.")

        private_key = rsa.generate_private_key(public_exponent=65537, key_size=self.KEY_SIZE, backend=default_backend())
        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)

        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:
        if not self.private_key_path.exists():
            raise FileNotFoundError(f"Private key not found at {self.private_key_path}.")
        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:
        if not self.public_key_path.exists():
            raise FileNotFoundError(f"Public key not found at {self.public_key_path}.")
        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:
        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:
        manifest_copy = _strip_signature_keys(dict(manifest))
        manifest_json = json.dumps(manifest_copy, sort_keys=True, separators=(",", ":"))
        manifest_bytes = manifest_json.encode("utf-8")
        private_key = self.load_private_key()
        signature = private_key.sign(
            manifest_bytes,
            padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
            hashes.SHA256(),
        )
        return base64.b64encode(signature).decode("utf-8")

    def verify_signature(
        self,
        manifest: Dict[str, Any],
        signature: str,
        public_key_pem: Optional[str] = None,
    ) -> bool:
        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:
        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:
        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
