nucleic-remote-ios/NucleicRemote/NucleicRemote/Models/IdentityStore.swift

import Foundation
import Security
import NucleicProtocol
#if canImport(UIKit)
import UIKit
#endif

/// What the phone pins about its Mac at pairing (SYNC §4.2): the host's static key (for IK
/// reconnect), a display name, and the transport + connection hint from the QR — LAN
/// host:port or the Mac's tailnet IP. The pairing secret is *not* stored — it's one-time.
/// Non-secret, so UserDefaults is fine; the device private key goes to Keychain. The new
/// optional fields decode as nil from a pre-transport record (= LAN).
struct PairedHost: Codable, Equatable {
    var deviceID: String
    var hostName: String
    var hostStaticKey: Data
    var fingerprint: String
    var lanHost: String?
    var lanPort: UInt16?
    /// `SyncTransportHint` raw value; nil = LAN (records saved before transports existed).
    var transport: String?
    var tailnetHost: String?
    var tailnetPort: UInt16?
    /// Nucleic Private Relay bootstrap (mesh P2): the host's room, this device's membership
    /// token, and an optional base-URL override. Set from the pairing QR and refreshed by
    /// the host's `relayMembership` push on every connect; nil while the host has the relay
    /// method off (records predating the relay decode them as nil).
    var relayRoomID: String?
    var relayMembershipToken: String?
    var relayURL: String?

    var transportHint: SyncTransportHint { transport.flatMap(SyncTransportHint.init(rawValue:)) ?? .lan }
}

/// Loads/persists this device's long-term `DeviceIdentity` (Keychain) and the pinned host
/// (UserDefaults). The identity is generated once on first launch and reused thereafter.
enum IdentityStore {
    private static let keychainAccount = "xyz.blakeslee.nucleic.remote.identity"
    /// Legacy single-host slot (pre-mesh P3). Migrated into `pairedHostsKey` on first registry read.
    private static let pairedHostKey = "nucleic.pairedHost"
    /// The multi-host registry (mesh P3): an ordered `[PairedHost]`, most-recently-paired last.
    /// The last entry is the "active" host the single connection uses today; the multiplexer will
    /// connect to all of them.
    private static let pairedHostsKey = "nucleic.pairedHosts"
    private static let deviceIDKey = "nucleic.deviceID"

    static func loadOrCreateIdentity() -> DeviceIdentity {
        if let data = keychainRead(), let identity = try? DeviceIdentity(importingRaw: data) {
            return identity
        }
        let identity = DeviceIdentity()
        keychainWrite(identity.exportRaw())
        return identity
    }

    /// Stable per-install device id (re-used across reconnects so the host can match the pin).
    static func deviceID() -> String {
        let defaults = UserDefaults.standard
        if let existing = defaults.string(forKey: deviceIDKey) { return existing }
        let id = deviceIDPrefix + UUID().uuidString.prefix(8).lowercased()
        defaults.set(id, forKey: deviceIDKey)
        return id
    }

    /// Idiom-tagged prefix so the host lists a paired device with the right kind/icon
    /// (`ipad-…` vs `iphone-…`). Only stamps *freshly generated* ids — an existing install
    /// keeps whatever id it already persisted, so upgrading a phone never changes its identity.
    private static var deviceIDPrefix: String {
        #if canImport(UIKit)
        return UIDevice.current.userInterfaceIdiom == .pad ? "ipad-" : "iphone-"
        #else
        return "iphone-"
        #endif
    }

    // MARK: - Paired-host registry (mesh P3)

    /// Every Mac this phone is paired with, most-recently-paired last. Migrates the legacy
    /// single-host slot into the registry on first read (then removes it), so an upgrade keeps its
    /// Mac. Ordered so the last is the "active" host today; a future multiplexer connects to all.
    static func pairedHosts() -> [PairedHost] {
        let defaults = UserDefaults.standard
        if let data = defaults.data(forKey: pairedHostsKey),
           let hosts = try? JSONDecoder().decode([PairedHost].self, from: data) {
            return hosts
        }
        // One-time migration from the pre-mesh single slot.
        if let legacy = defaults.data(forKey: pairedHostKey),
           let host = try? JSONDecoder().decode(PairedHost.self, from: legacy) {
            savePairedHosts([host])
            defaults.removeObject(forKey: pairedHostKey)
            return [host]
        }
        return []
    }

    /// The one paired Mac with `hostID` (its `fingerprint`), if any.
    static func pairedHost(id hostID: String) -> PairedHost? {
        pairedHosts().first { $0.fingerprint == hostID }
    }

    /// Add or update a host by fingerprint, moving it to the end (making it the active host). Used
    /// on a successful pairing handshake and on any address refresh.
    static func upsertPairedHost(_ host: PairedHost) {
        var hosts = pairedHosts().filter { $0.fingerprint != host.fingerprint }
        hosts.append(host)
        savePairedHosts(hosts)
    }

    /// Forget one paired Mac by fingerprint (per-host unpair).
    static func removePairedHost(id hostID: String) {
        savePairedHosts(pairedHosts().filter { $0.fingerprint != hostID })
    }

    /// Merge a Mac learned via mesh roster gossip (mesh "join"): append it if new, else refresh its
    /// hints **in place** — unlike `upsertPairedHost`, this must NOT reorder (which would hijack the
    /// active host) and must not clobber a relay token or our established device id with the
    /// gossiped nils.
    static func mergePairedHost(_ host: PairedHost) {
        var hosts = pairedHosts()
        if let index = hosts.firstIndex(where: { $0.fingerprint == host.fingerprint }) {
            var merged = host
            let existing = hosts[index]
            merged.deviceID = existing.deviceID  // keep our established id for this host
            merged.relayMembershipToken = host.relayMembershipToken ?? existing.relayMembershipToken
            merged.relayRoomID = host.relayRoomID ?? existing.relayRoomID
            merged.relayURL = host.relayURL ?? existing.relayURL
            hosts[index] = merged
        } else {
            hosts.append(host)
        }
        savePairedHosts(hosts)
    }

    /// Mutate one host's record **in place**, preserving registry order — unlike
    /// `upsertPairedHost`, this must not make the host active (it backs background
    /// refreshes like the relay-membership push, which can arrive from a non-active Mac).
    static func updatePairedHost(id hostID: String, mutate: (inout PairedHost) -> Void) {
        var hosts = pairedHosts()
        guard let index = hosts.firstIndex(where: { $0.fingerprint == hostID }) else { return }
        mutate(&hosts[index])
        savePairedHosts(hosts)
    }

    private static func savePairedHosts(_ hosts: [PairedHost]) {
        if let data = try? JSONEncoder().encode(hosts) {
            UserDefaults.standard.set(data, forKey: pairedHostsKey)
        }
    }

    // MARK: - Single-host bridge (until the RemoteStore multiplexer lands)

    /// The active host the single connection uses — the most-recently-paired. `nil` if unpaired.
    static func loadPairedHost() -> PairedHost? {
        pairedHosts().last
    }

    /// Pair (or re-pair) a host and make it active.
    static func savePairedHost(_ host: PairedHost) {
        upsertPairedHost(host)
    }

    /// Unpair the active host (the one the single connection currently uses).
    static func clearPairedHost() {
        if let active = pairedHosts().last { removePairedHost(id: active.fingerprint) }
    }

    // MARK: - Keychain

    private static func keychainRead() -> Data? {
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrAccount as String: keychainAccount,
            kSecReturnData as String: true,
            kSecMatchLimit as String: kSecMatchLimitOne,
        ]
        var item: CFTypeRef?
        guard SecItemCopyMatching(query as CFDictionary, &item) == errSecSuccess else { return nil }
        return item as? Data
    }

    private static func keychainWrite(_ data: Data) {
        let delete: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrAccount as String: keychainAccount,
        ]
        SecItemDelete(delete as CFDictionary)
        let add: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrAccount as String: keychainAccount,
            kSecValueData as String: data,
            kSecAttrAccessible as String: kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly,
        ]
        SecItemAdd(add as CFDictionary, nil)
    }
}