import Foundation
import os
import Security

/// Persists the device identity so it survives relaunch. The raw `.p12` bytes
/// **and** its passphrase are stored together (the passphrase is required to
/// re-import via `SecPKCS12Import` at every launch), then re-imported on load.
public protocol ClientIdentityStore: Sendable {
    /// Validate (`SecPKCS12Import`) then persist the `.p12` + passphrase.
    /// Throws `PKCS12ImportError` on a bad passphrase / corrupt file (nothing is
    /// persisted in that case) and `ClientIdentityStoreError` on a storage fault.
    func save(p12Data: Data, passphrase: String) throws
    /// Re-import and return the stored identity; `nil` if none is installed.
    func loadIdentity() throws -> ClientIdentity?
    /// Display summary of the stored certificate; `nil` if none is installed.
    func loadSummary() throws -> ClientCertificateSummary?
    /// Delete the stored identity (rotation / removal). Idempotent.
    func remove() throws
    /// Cheap existence check for gating (does NOT re-import).
    func hasInstalledIdentity() -> Bool
}

public enum ClientIdentityStoreError: Error, Equatable, Sendable {
    /// A Keychain `SecItem*` call failed with this `OSStatus`.
    case keychain(OSStatus)
    /// The stored blob was present but could not be decoded.
    case corruptStoredBlob
}

public extension ClientIdentityStore {
    /// Convenience for composition roots: load the identity, logging and
    /// swallowing errors into `nil`. A missing cert is the normal pre-install
    /// state; a genuine fault must not crash launch, but is logged (never
    /// silently dropped).
    func loadedIdentityOrNil() -> ClientIdentity? {
        do {
            return try loadIdentity()
        } catch {
            ClientTLSLog.identity.error(
                "loadIdentity failed: \(String(describing: error), privacy: .public)"
            )
            return nil
        }
    }
}

/// The stored payload — `.p12` bytes plus the passphrase needed to re-import.
private struct StoredP12Blob: Codable {
    let p12: Data
    let passphrase: String
}

/// Keychain-backed store: one `kSecClassGenericPassword` item holding the
/// JSON-encoded `StoredP12Blob` in `kSecValueData`, protected with
/// `kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly` (available after the first
/// unlock post-boot, never migrates off this device).
public struct KeychainClientIdentityStore: ClientIdentityStore {
    public static let defaultService = "com.yaojia.webterm.clienttls"
    public static let defaultAccount = "device-identity"

    private let service: String
    private let account: String

    public init(
        service: String = defaultService, account: String = defaultAccount
    ) {
        self.service = service
        self.account = account
    }

    public func save(p12Data: Data, passphrase: String) throws {
        // Validate BEFORE persisting — a wrong passphrase / corrupt file must
        // surface to the install UI and leave any prior identity untouched.
        _ = try PKCS12Importer.importIdentity(data: p12Data, passphrase: passphrase)
        let blob = try encode(StoredP12Blob(p12: p12Data, passphrase: passphrase))
        try writeItem(blob)
    }

    public func loadIdentity() throws -> ClientIdentity? {
        guard let blob = try readBlob() else { return nil }
        return try PKCS12Importer.importIdentity(
            data: blob.p12, passphrase: blob.passphrase
        )
    }

    public func loadSummary() throws -> ClientCertificateSummary? {
        try loadIdentity()?.summary()
    }

    public func remove() throws {
        let status = SecItemDelete(baseQuery() as CFDictionary)
        guard status == errSecSuccess || status == errSecItemNotFound else {
            throw ClientIdentityStoreError.keychain(status)
        }
    }

    public func hasInstalledIdentity() -> Bool {
        var query = baseQuery()
        query[kSecReturnData as String] = false
        query[kSecMatchLimit as String] = kSecMatchLimitOne
        return SecItemCopyMatching(query as CFDictionary, nil) == errSecSuccess
    }

    // MARK: - Keychain plumbing

    private func baseQuery() -> [String: Any] {
        [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: service,
            kSecAttrAccount as String: account,
        ]
    }

    private func writeItem(_ data: Data) throws {
        // Delete-then-add keeps the item's protection class deterministic
        // (SecItemUpdate can't change kSecAttrAccessible in place).
        let deleteStatus = SecItemDelete(baseQuery() as CFDictionary)
        guard deleteStatus == errSecSuccess || deleteStatus == errSecItemNotFound else {
            throw ClientIdentityStoreError.keychain(deleteStatus)
        }
        var attributes = baseQuery()
        attributes[kSecValueData as String] = data
        attributes[kSecAttrAccessible as String] =
            kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly
        let addStatus = SecItemAdd(attributes as CFDictionary, nil)
        guard addStatus == errSecSuccess else {
            throw ClientIdentityStoreError.keychain(addStatus)
        }
    }

    private func readBlob() throws -> StoredP12Blob? {
        var query = baseQuery()
        query[kSecReturnData as String] = true
        query[kSecMatchLimit as String] = kSecMatchLimitOne
        var result: CFTypeRef?
        let status = SecItemCopyMatching(query as CFDictionary, &result)
        if status == errSecItemNotFound { return nil }
        guard status == errSecSuccess, let data = result as? Data else {
            throw ClientIdentityStoreError.keychain(status)
        }
        do {
            return try JSONDecoder().decode(StoredP12Blob.self, from: data)
        } catch {
            throw ClientIdentityStoreError.corruptStoredBlob
        }
    }

    private func encode(_ blob: StoredP12Blob) throws -> Data {
        do {
            return try JSONEncoder().encode(blob)
        } catch {
            throw ClientIdentityStoreError.corruptStoredBlob
        }
    }
}

/// In-memory store for previews and unit tests: same import/summary code path as
/// the keychain store (so the roundtrip is exercised) without any keychain
/// entitlement. `@unchecked Sendable` — mutable blob guarded by a lock.
public final class InMemoryClientIdentityStore: ClientIdentityStore, @unchecked Sendable {
    private let lock = NSLock()
    private var blob: StoredP12BlobBox?

    /// Boxed so the private `StoredP12Blob` type stays file-private above; this
    /// mirror keeps the two bytes+passphrase without exposing the Codable type.
    private struct StoredP12BlobBox {
        let p12: Data
        let passphrase: String
    }

    public init() {}

    public func save(p12Data: Data, passphrase: String) throws {
        _ = try PKCS12Importer.importIdentity(data: p12Data, passphrase: passphrase)
        lock.withLock { blob = StoredP12BlobBox(p12: p12Data, passphrase: passphrase) }
    }

    public func loadIdentity() throws -> ClientIdentity? {
        guard let stored = lock.withLock({ blob }) else { return nil }
        return try PKCS12Importer.importIdentity(
            data: stored.p12, passphrase: stored.passphrase
        )
    }

    public func loadSummary() throws -> ClientCertificateSummary? {
        try loadIdentity()?.summary()
    }

    public func remove() throws {
        lock.withLock { blob = nil }
    }

    public func hasInstalledIdentity() -> Bool {
        lock.withLock { blob != nil }
    }
}

enum ClientTLSLog {
    static let identity = Logger(subsystem: "com.yaojia.webterm", category: "client-tls")
}