"""PKI operations: internal CA + server certificate issuance (spec 16-19).

All private-key and signing operations happen here, in the backend —
the frontend only triggers them (spec section 16). Keys are returned in
PEM form to the storage layer, which encrypts them at rest (spec 43).

Release 2 MVP scope:
  * create an internal CA (RSA 3072/4096, ECDSA P-256/P-384)
  * issue a server certificate from that CA (CN + DNS/IP SANs, serverAuth)
CSR signing, client certificates, revocation/CRL arrive later.
"""

from __future__ import annotations

import datetime
import uuid
from dataclasses import dataclass

from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec, ed448, ed25519, rsa
from cryptography.x509.oid import NameOID

# Key algorithms supported for both CAs and leaf certificates (spec 17).
KEY_ALGORITHMS = ("rsa3072", "rsa4096", "ecdsa-p256", "ecdsa-p384")


class PkiError(ValueError):
    """Invalid PKI request (bad algorithm, empty subject, ...)."""


def _generate_key(algorithm: str):
    """Generate a private key for the given algorithm identifier."""
    if algorithm == "rsa3072":
        return rsa.generate_private_key(public_exponent=65537, key_size=3072)
    if algorithm == "rsa4096":
        return rsa.generate_private_key(public_exponent=65537, key_size=4096)
    if algorithm == "ecdsa-p256":
        return ec.generate_private_key(ec.SECP256R1())
    if algorithm == "ecdsa-p384":
        return ec.generate_private_key(ec.SECP384R1())
    raise PkiError(
        f"unsupported key algorithm {algorithm!r}; expected one of {KEY_ALGORITHMS}"
    )


def _name(common_name: str, organization: str | None, unit: str | None,
          country: str | None) -> x509.Name:
    attrs = [x509.NameAttribute(NameOID.COMMON_NAME, common_name)]
    if organization:
        attrs.append(x509.NameAttribute(NameOID.ORGANIZATION_NAME, organization))
    if unit:
        attrs.append(x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, unit))
    if country:
        attrs.append(x509.NameAttribute(NameOID.COUNTRY_NAME, country))
    return x509.Name(attrs)


def _pem_key(key) -> str:
    return key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    ).decode("utf-8")


def _pem_cert(cert: x509.Certificate) -> str:
    return cert.public_bytes(serialization.Encoding.PEM).decode("utf-8")


def _serial() -> int:
    # Positive, 159-bit (20-byte) random serial per RFC 5280 conventions.
    return int.from_bytes(uuid.uuid4().bytes[:20], "big") >> 1


def _not_after(now: datetime.datetime, validity_days: int) -> datetime.datetime:
    if validity_days <= 0:
        raise PkiError("validity must be a positive number of days")
    return now + datetime.timedelta(days=validity_days)


@dataclass
class IssuedPair:
    """A certificate + its private key, both as PEM text."""

    certificate_pem: str
    private_key_pem: str
    serial_number: int
    not_valid_after: datetime.datetime


def create_ca(
    *,
    common_name: str,
    validity_days: int,
    key_algorithm: str,
    organization: str | None = None,
    organizational_unit: str | None = None,
    country: str | None = None,
) -> IssuedPair:
    """Create a self-signed CA certificate (spec section 17)."""
    if not common_name.strip():
        raise PkiError("common name must not be empty")
    now = datetime.datetime.now(datetime.UTC)
    key = _generate_key(key_algorithm)
    subject = issuer = _name(common_name, organization, organizational_unit, country)
    cert = (
        x509.CertificateBuilder()
        .subject_name(subject)
        .issuer_name(issuer)
        .public_key(key.public_key())
        .serial_number(_serial())
        .not_valid_before(now - datetime.timedelta(minutes=5))  # clock skew
        .not_valid_after(_not_after(now, validity_days))
        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
        .add_extension(x509.KeyUsage(
            digital_signature=True,
            content_commitment=False,
            key_encipherment=False,
            data_encipherment=False,
            key_agreement=False,
            key_cert_sign=True,
            crl_sign=True,
            encipher_only=False,
            decipher_only=False,
        ), critical=True)
        .sign(key, hashes.SHA256())
    )
    return IssuedPair(
        certificate_pem=_pem_cert(cert),
        private_key_pem=_pem_key(key),
        serial_number=cert.serial_number,
        not_valid_after=cert.not_valid_after_utc,
    )


def issue_server_certificate(
    *,
    ca_certificate_pem: str,
    ca_private_key_pem: str,
    common_name: str,
    validity_days: int,
    dns_sans: list[str] | None = None,
    ip_sans: list[str] | None = None,
    organization: str | None = None,
    organizational_unit: str | None = None,
    country: str | None = None,
) -> IssuedPair:
    """Issue a server certificate signed by the given CA (spec section 19).

    The certificate carries the serverAuth extended key usage and the
    CN plus any DNS/IP SANs — exactly what a Mosquitto TLS listener needs.
    """
    if not common_name.strip():
        raise PkiError("common name must not be empty")
    dns_sans = [s.strip() for s in (dns_sans or []) if s.strip()]
    ip_sans = [s.strip() for s in (ip_sans or []) if s.strip()]

    ca_cert = x509.load_pem_x509_certificate(ca_certificate_pem.encode("utf-8"))
    ca_key = serialization.load_pem_private_key(
        ca_private_key_pem.encode("utf-8"), password=None
    )
    # Type narrowing for the checker: certificate signing requires an
    # RSA/EC/Ed25519/Ed448 key — DH/DSA keys load but cannot sign certs.
    if not isinstance(
        ca_key,
        (
            rsa.RSAPrivateKey,
            ec.EllipticCurvePrivateKey,
            ed25519.Ed25519PrivateKey,
            ed448.Ed448PrivateKey,
        ),
    ):
        raise PkiError("the CA private key type cannot sign certificates")

    # The leaf key is always ECDSA P-256 for Release 2: fast to generate,
    # modern, and universally supported by MQTT clients.
    key = _generate_key("ecdsa-p256")
    now = datetime.datetime.now(datetime.UTC)
    subject = _name(common_name, organization, organizational_unit, country)

    san_entries: list[x509.GeneralName] = [x509.DNSName(common_name)]
    for name in dns_sans:
        san_entries.append(x509.DNSName(name))
    for ip in ip_sans:
        try:
            import ipaddress

            san_entries.append(x509.IPAddress(ipaddress.ip_address(ip)))
        except ValueError as exc:
            raise PkiError(f"invalid IP SAN {ip!r}") from exc

    cert = (
        x509.CertificateBuilder()
        .subject_name(subject)
        .issuer_name(ca_cert.subject)
        .public_key(key.public_key())
        .serial_number(_serial())
        .not_valid_before(now - datetime.timedelta(minutes=5))
        .not_valid_after(_not_after(now, validity_days))
        .add_extension(
            x509.SubjectAlternativeName(san_entries), critical=False
        )
        .add_extension(
            x509.ExtendedKeyUsage([x509.oid.ExtendedKeyUsageOID.SERVER_AUTH]),
            critical=False,
        )
        .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
        .sign(ca_key, hashes.SHA256())
    )
    return IssuedPair(
        certificate_pem=_pem_cert(cert),
        private_key_pem=_pem_key(key),
        serial_number=cert.serial_number,
        not_valid_after=cert.not_valid_after_utc,
    )


def fingerprint(certificate_pem: str) -> str:
    """SHA-256 fingerprint of a certificate, colon-separated hex."""
    cert = x509.load_pem_x509_certificate(certificate_pem.encode("utf-8"))
    digest = cert.fingerprint(hashes.SHA256())
    return ":".join(f"{b:02x}" for b in digest)