import Crypto
import Foundation
import SwiftASN1
import X509
import _CryptoExtras

/// The local certificate authority behind the Nucleic MDM plane (docs/MACOS_VM_MDM.md §4).
///
/// It issues three kinds of certificate, all chaining to one self-signed root:
///   • the **root** itself — staged into the guest and marked `trustRoot` before enrollment so the
///     MDM server's TLS and the SCEP handshake validate;
///   • the **server TLS identity** — a leaf with `serverAuth` EKU and a stable DNS SAN
///     (``MDMConfiguration/serverDNSName``) that the guest pins to the vmnet gateway;
///   • the **device identity** — a leaf issued from the SCEP CSR `mdmclient` presents at enrollment
///     (``issueDeviceCertificate(fromCSR:)``), which then signs every check-in/command request
///     (`Mdm-Signature`).
///
/// The CA is **RSA** end to end: mdmclient's SCEP client and the `Mdm-Signature`/TLS paths all
/// interoperate with RSA + SHA-256, and the CA private key doubles as the SCEP decryption key
/// (`_RSA.Encryption`), so keeping one algorithm avoids a second keypair.
///
/// Pure and host-only in effect (no Network / Virtualization): the CA is created once per base
/// build, persisted next to the base bundle (PEM), and reloaded for Mode B. Nothing here touches
/// the Keychain — the private key lives in a host-owned file, never staged into any guest.
public struct MDMCertificateAuthority: Sendable {
    /// The self-signed root certificate (also the SCEP CA cert served at `GetCACert`).
    public let rootCertificate: Certificate
    /// The root's RSA private key, used to sign issued leaves and to decrypt SCEP `PKIOperation`
    /// envelopes. Never leaves the host.
    private let rootKey: _RSA.Signing.PrivateKey
    /// The distinguished name shared by the root subject/issuer.
    public let caName: DistinguishedName

    /// A monotonic serial source. The root is serial 1; issued leaves take 2, 3, … Persisted with
    /// the CA so a reload keeps issuing unique serials.
    private let nextSerial: Int

    // MARK: - Creation

    /// Errors from CA operations.
    public enum CAError: Error, CustomStringConvertible {
        case malformedPEM(String)
        case csrRejected(String)
        case issuance(String)

        public var description: String {
            switch self {
            case .malformedPEM(let s): return "malformed CA PEM material: \(s)"
            case .csrRejected(let s): return "SCEP CSR rejected: \(s)"
            case .issuance(let s): return "certificate issuance failed: \(s)"
            }
        }
    }

    /// Generate a brand-new CA (fresh RSA-2048 root). `commonName` names both the root and the org
    /// the enrollment profile advertises. `validityYears` bounds the root (leaves are bounded
    /// separately and always within the root's window).
    public static func generate(
        commonName: String = "Nucleic MDM Root CA",
        organization: String = "Nucleic",
        validityYears: Int = 10,
        now: Date = Date()
    ) throws -> MDMCertificateAuthority {
        let key = try _RSA.Signing.PrivateKey(keySize: .bits2048)
        let name = try DistinguishedName {
            CommonName(commonName)
            OrganizationName(organization)
        }
        let certKey = Certificate.PrivateKey(key)
        let notBefore = now.addingTimeInterval(-3600)  // tolerate guest clock skew
        let notAfter = now.addingTimeInterval(Double(validityYears) * 365 * 86_400)
        let extensions = try Certificate.Extensions {
            Critical(BasicConstraints.isCertificateAuthority(maxPathLength: 1))
            Critical(KeyUsage(keyCertSign: true, cRLSign: true))
            SubjectKeyIdentifier(hash: certKey.publicKey)
        }
        let root = try Certificate(
            version: .v3,
            serialNumber: Certificate.SerialNumber(bytes: [1]),
            publicKey: certKey.publicKey,
            notValidBefore: notBefore,
            notValidAfter: notAfter,
            issuer: name,
            subject: name,
            signatureAlgorithm: .sha256WithRSAEncryption,
            extensions: extensions,
            issuerPrivateKey: certKey)
        return MDMCertificateAuthority(
            rootCertificate: root, rootKey: key, caName: name, nextSerial: 2)
    }

    private init(
        rootCertificate: Certificate, rootKey: _RSA.Signing.PrivateKey,
        caName: DistinguishedName, nextSerial: Int
    ) {
        self.rootCertificate = rootCertificate
        self.rootKey = rootKey
        self.caName = caName
        self.nextSerial = nextSerial
    }

    // MARK: - Persistence (PEM)

    /// The on-disk form: the root cert PEM, the root key PEM, and the next serial. Kept host-side
    /// next to the base bundle; the private-key PEM is the crown jewel (see §8 trust boundary).
    public struct Persisted: Codable, Sendable {
        public var rootCertPEM: String
        public var rootKeyPEM: String
        public var nextSerial: Int
        public init(rootCertPEM: String, rootKeyPEM: String, nextSerial: Int) {
            self.rootCertPEM = rootCertPEM
            self.rootKeyPEM = rootKeyPEM
            self.nextSerial = nextSerial
        }
    }

    /// Serialize to the persisted PEM form.
    public func persisted() throws -> Persisted {
        Persisted(
            rootCertPEM: try rootCertificate.serializeAsPEM().pemString,
            rootKeyPEM: try Certificate.PrivateKey(rootKey).serializeAsPEM().pemString,
            nextSerial: nextSerial)
    }

    /// Reload a CA from its persisted PEM form.
    public init(persisted p: Persisted) throws {
        let cert: Certificate
        do { cert = try Certificate(pemEncoded: p.rootCertPEM) }
        catch { throw CAError.malformedPEM("root cert: \(error)") }
        let key: _RSA.Signing.PrivateKey
        do {
            let doc = try PEMDocument(pemString: p.rootKeyPEM)
            key = try _RSA.Signing.PrivateKey(derRepresentation: doc.derBytes)
        } catch { throw CAError.malformedPEM("root key: \(error)") }
        self.init(
            rootCertificate: cert, rootKey: key, caName: cert.subject, nextSerial: p.nextSerial)
    }

    // MARK: - Issuance

    /// The root as DER — the `.cer` staged into the guest for `security add-trusted-cert` and the
    /// SCEP `GetCACert` reply.
    public func rootCertificateDER() throws -> [UInt8] {
        var s = DER.Serializer()
        try rootCertificate.serialize(into: &s)
        return s.serializedBytes
    }

    /// Issue the **server TLS identity**: a leaf keyed to `serverDNSName` (+ any extra SANs, e.g. a
    /// gateway IP) with `serverAuth` EKU. Returns the leaf plus its freshly generated private key.
    ///
    /// `validityDays` defaults to **397** — Apple's `SecTrust` TLS policy rejects a server cert whose
    /// lifetime exceeds ~398 days ("certificate exceeds maximum temporal validity period"), so a base
    /// used for Mode B beyond that window must re-issue the server identity (cheap; the CA persists).
    public func issueServerCertificate(
        dnsName: String, ipSANs: [String] = [], validityDays: Int = 397, now: Date = Date()
    ) throws -> (certificate: Certificate, privateKey: _RSA.Signing.PrivateKey) {
        let leafKey = try _RSA.Signing.PrivateKey(keySize: .bits2048)
        let leafCertKey = Certificate.PrivateKey(leafKey)
        var sans: [GeneralName] = [.dnsName(dnsName)]
        for ip in ipSANs {
            if let name = MDMCertificateAuthority.ipSAN(ip) { sans.append(name) }
        }
        let subject = try DistinguishedName { CommonName(dnsName) }
        let extensions = try Certificate.Extensions {
            Critical(BasicConstraints.notCertificateAuthority)
            KeyUsage(digitalSignature: true, keyEncipherment: true)
            try ExtendedKeyUsage([.serverAuth])
            SubjectAlternativeNames(sans)
            SubjectKeyIdentifier(hash: leafCertKey.publicKey)
            AuthorityKeyIdentifier(keyIdentifier: rootSubjectKeyID())
        }
        let cert = try signLeaf(
            subject: subject, publicKey: leafCertKey.publicKey, extensions: extensions,
            validityDays: validityDays, now: now)
        return (cert, leafKey)
    }

    /// Issue a leaf for **signing the enrollment profile** (CMS). Distinct from the root: a CA cert's
    /// KeyUsage is `keyCertSign`/`cRLSign`, so signing content with it violates its own constraints —
    /// this leaf carries `digitalSignature`. The guest verifies the signature against the CA root it
    /// already trusts (`add-trusted-cert`), which turns the profile from "Not Signed" into a verified
    /// one whose certificate payloads macOS is willing to trust.
    public func issueProfileSigningCertificate(
        commonName: String = "Nucleic Profile Signing", validityDays: Int = 1825, now: Date = Date()
    ) throws -> (certificate: Certificate, privateKey: _RSA.Signing.PrivateKey) {
        let leafKey = try _RSA.Signing.PrivateKey(keySize: .bits2048)
        let leafCertKey = Certificate.PrivateKey(leafKey)
        let subject = try DistinguishedName { CommonName(commonName) }
        let extensions = try Certificate.Extensions {
            Critical(BasicConstraints.notCertificateAuthority)
            KeyUsage(digitalSignature: true)
            SubjectKeyIdentifier(hash: leafCertKey.publicKey)
            AuthorityKeyIdentifier(keyIdentifier: rootSubjectKeyID())
        }
        let cert = try signLeaf(
            subject: subject, publicKey: leafCertKey.publicKey, extensions: extensions,
            validityDays: validityDays, now: now)
        return (cert, leafKey)
    }

    /// Issue the **device identity** from a CSR (host-generated for the baked PKCS#12 today; the
    /// deferred SCEP path would feed the device's own CSR here). The CSR's public key and subject are
    /// honored; we stamp `digitalSignature`/`keyEncipherment` KeyUsage and `clientAuth` EKU so the cert
    /// can sign `Mdm-Signature`. One shared logical device (§2.5), so the subject need not be unique.
    /// The device cert is a client cert (not TLS server), so the 398-day cap does not apply — 5 years.
    public func issueDeviceCertificate(
        fromCSR csr: CertificateSigningRequest, validityDays: Int = 1825, now: Date = Date()
    ) throws -> Certificate {
        let subject = csr.subject.isEmpty
            ? (try DistinguishedName { CommonName("Nucleic Managed Device") })
            : csr.subject
        let extensions = try Certificate.Extensions {
            Critical(BasicConstraints.notCertificateAuthority)
            KeyUsage(digitalSignature: true, keyEncipherment: true)
            try ExtendedKeyUsage([.clientAuth])
            SubjectKeyIdentifier(hash: csr.publicKey)
            AuthorityKeyIdentifier(keyIdentifier: rootSubjectKeyID())
        }
        return try signLeaf(
            subject: subject, publicKey: csr.publicKey, extensions: extensions,
            validityDays: validityDays, now: now)
    }

    // MARK: - Internals

    private func signLeaf(
        subject: DistinguishedName, publicKey: Certificate.PublicKey,
        extensions: Certificate.Extensions, validityDays: Int, now: Date
    ) throws -> Certificate {
        let notBefore = now.addingTimeInterval(-3600)
        let notAfter = min(
            now.addingTimeInterval(Double(validityDays) * 86_400),
            rootCertificate.notValidAfter)
        do {
            return try Certificate(
                version: .v3,
                serialNumber: Certificate.SerialNumber(bytes: ArraySlice(Self.randomSerialBytes())),
                publicKey: publicKey,
                notValidBefore: notBefore,
                notValidAfter: notAfter,
                issuer: caName,
                subject: subject,
                signatureAlgorithm: .sha256WithRSAEncryption,
                extensions: extensions,
                issuerPrivateKey: Certificate.PrivateKey(rootKey))
        } catch {
            throw CAError.issuance("\(error)")
        }
    }

    /// The root's Subject Key Identifier, recomputed from its public key (the same value we stamped
    /// into the root at generation) — used as issued leaves' Authority Key Identifier.
    private func rootSubjectKeyID() -> ArraySlice<UInt8> {
        SubjectKeyIdentifier(hash: rootCertificate.publicKey).keyIdentifier
    }

    /// A fresh random 16-byte positive serial for each issued leaf.
    ///
    /// Serials MUST be unique per issuer: X.509 identifies a certificate by (issuer, serialNumber),
    /// and Apple's keychain/trust machinery keys certificate lookups on exactly that pair. The old
    /// counter could not work — the CA is a value type, so the `let nextSerial` never advanced across
    /// non-mutating issuance and *every* leaf was stamped serial 2. The device identity baked into the
    /// enrollment profile and the MDM server's TLS leaf therefore collided, and a live macOS 27 run
    /// failed profile installation with "The certificate could not be verified (authentication error)".
    /// Randomness also keeps serials unique across CA reloads, which a persisted counter would not
    /// guarantee. 16 bytes satisfies RFC 5280 §4.1.2.2 (≤20 octets) and the CA/B ≥64-bit entropy rule;
    /// the top bit is cleared so the ASN.1 INTEGER is positive, and a zero leading octet is avoided so
    /// the encoding stays minimal.
    static func randomSerialBytes() -> [UInt8] {
        var bytes = [UInt8](repeating: 0, count: 16)
        for i in bytes.indices { bytes[i] = UInt8.random(in: 0...255) }
        bytes[0] &= 0x7F
        if bytes[0] == 0 { bytes[0] = 1 }
        return bytes
    }

    /// Parse a dotted-quad or IPv6 literal into a SAN `ipAddress` octet string, or `nil` if unparseable.
    static func ipSAN(_ ip: String) -> GeneralName? {
        let parts = ip.split(separator: ".")
        if parts.count == 4, let octets = try? parts.map({ part -> UInt8 in
            guard let n = UInt8(part) else { throw CAError.issuance("bad octet") }
            return n
        }) {
            return .ipAddress(ASN1OctetString(contentBytes: ArraySlice(octets)))
        }
        return nil
    }
}
