RouterOS/RouterOSAssistant/Core/Services/ConnectionService.swift

import Foundation

/// Drives the REST-first, SSH-fallback connection flow and holds the connected device's state.
@MainActor
final class ConnectionService: ObservableObject {
    enum State: Equatable {
        case idle
        case connecting
        case connected(kind: RouterOSTransportKind)
        case needsCertificateConfirmation(fingerprint: String)
        case needsSSHHostKeyConfirmation(fingerprint: String)
        case failed(String)
    }

    @Published private(set) var state: State = .idle
    @Published private(set) var deviceInfo: RouterDeviceInfo?
    @Published private(set) var routerBoardInfo: RouterBoardInfo?
    @Published private(set) var interfaces: [NetworkInterface] = []
    /// Set when a REST connection succeeds but the router's SSH host key is still untrusted —
    /// distinct from `state`'s `.needsSSHHostKeyConfirmation` (which blocks the initial connect
    /// attempt itself). Dedicated SSH-only services (`BackupService`, `UpdateService`,
    /// `NetworkToolsService`, `InterfaceTrafficMonitor`, `FactoryResetService`) always need SSH
    /// regardless of the active transport — a REST-first connection never establishes SSH trust
    /// on its own, so their first real use (e.g. the automatic pre-apply backup) used to dead-end
    /// with a raw "Unbekannter SSH-Schlüssel" error and no way to confirm it (Bug 13, see
    /// HANDOFF.md). Checked proactively right after every successful REST connect instead, so
    /// trust is already established by the time any of those services runs.
    @Published private(set) var pendingSSHTrustFingerprint: String?
    /// Credentials of the current (or last attempted) connection, shared with features
    /// that need their own dedicated connection, e.g. BackupService's SSH export.
    @Published private(set) var credentials: RouterOSCredentials?

    private let certificateTrust: CertificateTrustStore
    private let sshHostKeyTrust: SSHHostKeyTrustStore
    private var activeTransport: RouterOSTransport?
    /// Whether the Expert tool has already made its one automatic safety backup for this
    /// connection — it backs up before its first write, not before every single edit (unlike
    /// the Einrichten-Wizard, which backs up before every "Jetzt anwenden"). Reset on
    /// `disconnect()` so a new connection gets a fresh backup again.
    private(set) var hasExpertToolBackedUpThisSession = false

    func markExpertToolBackedUpThisSession() {
        hasExpertToolBackedUpThisSession = true
    }

    init(
        certificateTrust: CertificateTrustStore = CertificateTrustStore(),
        sshHostKeyTrust: SSHHostKeyTrustStore = SSHHostKeyTrustStore()
    ) {
        self.certificateTrust = certificateTrust
        self.sshHostKeyTrust = sshHostKeyTrust
    }

    /// Injection point for tests: bypasses the real REST/SSH transports.
    func connect(with credentials: RouterOSCredentials, makeRestTransport: () -> RouterOSTransport, makeSSHTransport: () -> RouterOSTransport) async {
        state = .connecting
        self.credentials = credentials

        let rest = makeRestTransport()
        do {
            try await rest.connect()
            await finishConnecting(using: rest)
            await verifySSHTrust(makeSSHTransport: makeSSHTransport)
            return
        } catch RouterOSError.untrustedCertificate(let fingerprint) {
            state = .needsCertificateConfirmation(fingerprint: fingerprint)
            return
        } catch {
            // REST fehlgeschlagen (z.B. altes RouterOS ohne REST-API) -> SSH-Fallback versuchen.
        }

        let ssh = makeSSHTransport()
        do {
            try await ssh.connect()
            await finishConnecting(using: ssh)
        } catch RouterOSError.untrustedSSHHostKey(let fingerprint) {
            state = .needsSSHHostKeyConfirmation(fingerprint: fingerprint)
        } catch {
            state = .failed(error.localizedDescription)
        }
    }

    func connect(with credentials: RouterOSCredentials) async {
        await connect(
            with: credentials,
            makeRestTransport: { RestTransport(credentials: credentials, certificateTrust: self.certificateTrust) },
            makeSSHTransport: { SSHTransport(credentials: credentials, hostKeyTrust: self.sshHostKeyTrust) }
        )
    }

    func trustCurrentCertificateAndRetry(fingerprint: String) async {
        guard let credentials else { return }
        certificateTrust.trust(host: credentials.host, fingerprint: fingerprint)
        await connect(with: credentials)
    }

    func trustCurrentSSHHostKeyAndRetry(fingerprint: String) async {
        guard let credentials else { return }
        sshHostKeyTrust.trust(host: credentials.host, fingerprint: fingerprint)
        await connect(with: credentials)
    }

    /// Confirms `pendingSSHTrustFingerprint` — unlike `trustCurrentSSHHostKeyAndRetry`, no
    /// reconnect needed: the active REST transport is unaffected, this only unblocks whichever
    /// dedicated SSH service runs next (see `pendingSSHTrustFingerprint`'s doc comment).
    func trustPendingSSHHostKeyAndRetry() {
        guard let credentials, let fingerprint = pendingSSHTrustFingerprint else { return }
        sshHostKeyTrust.trust(host: credentials.host, fingerprint: fingerprint)
        pendingSSHTrustFingerprint = nil
    }

    func dismissPendingSSHTrust() {
        pendingSSHTrustFingerprint = nil
    }

    /// Best-effort SSH host-key probe run once right after a successful REST connect — see
    /// `pendingSSHTrustFingerprint`'s doc comment for why this is needed at all. A throwaway SSH
    /// connection is the only way to learn the router's host key (there's no way to ask for it
    /// without actually connecting); closed immediately after, since this connection is never
    /// otherwise used. Any error other than `untrustedSSHHostKey` is swallowed deliberately — REST
    /// already connected successfully, so e.g. a wrong SSH port or a transient network hiccup here
    /// must not disrupt the working connection; it'll surface normally whenever a dedicated SSH
    /// service actually needs it.
    private func verifySSHTrust(makeSSHTransport: () -> RouterOSTransport) async {
        let ssh = makeSSHTransport()
        do {
            try await ssh.connect()
            await ssh.disconnect()
        } catch RouterOSError.untrustedSSHHostKey(let fingerprint) {
            pendingSSHTrustFingerprint = fingerprint
        } catch {
            // Swallowed — see doc comment above.
        }
    }

    /// Applies a single configuration change on the active transport (REST or SSH).
    func apply(_ command: RouterOSCommand) async throws {
        guard let activeTransport else { throw RouterOSError.notConnected }
        try await activeTransport.apply(command)
    }

    func fetchFirewallRuleCounts() async throws -> FirewallRuleCounts {
        guard let activeTransport else { throw RouterOSError.notConnected }
        return try await activeTransport.fetchFirewallRuleCounts()
    }

    /// Generic read for the Expert tool — lists existing items under any RouterOS menu path.
    func fetchMenuItems(menuPath: String, restPath: String) async throws -> [RouterOSMenuItem] {
        guard let activeTransport else { throw RouterOSError.notConnected }
        return try await activeTransport.fetchMenuItems(menuPath: menuPath, restPath: restPath)
    }

    /// Field values matching an internal-property filter — see `RouterOSTransport.fetchFieldValues`.
    func fetchFieldValues(menuPath: String, restPath: String, whereField: String, whereValue: String, returnField: String) async throws -> Set<String> {
        guard let activeTransport else { throw RouterOSError.notConnected }
        return try await activeTransport.fetchFieldValues(menuPath: menuPath, restPath: restPath, whereField: whereField, whereValue: whereValue, returnField: returnField)
    }

    /// Live, read-only check for the Setup wizard's LAN step: whether `interfaceName` already
    /// carries configuration that assigning it its own LAN/DHCP role would silently override
    /// (bridge membership, an existing IP address, or already being a WAN dial-up). "bridge"
    /// itself is never flagged — it's the app's own shared-LAN interface, never "someone else's"
    /// config. Never modifies anything; the caller decides what, if anything, to remove.
    func checkPortConflict(interfaceName: String) async throws -> PortConflict? {
        guard interfaceName != "bridge" else { return nil }

        async let bridgePorts = fetchMenuItems(menuPath: "/interface bridge port", restPath: "interface/bridge/port")
        async let addresses = fetchMenuItems(menuPath: "/ip address", restPath: "ip/address")
        async let dhcpClients = fetchMenuItems(menuPath: "/ip dhcp-client", restPath: "ip/dhcp-client")
        async let pppoeClients = fetchMenuItems(menuPath: "/interface pppoe-client", restPath: "interface/pppoe-client")

        var reasons: [PortConflict.Reason] = []
        if let bridgeName = try await bridgePorts.first(where: { $0.fields["interface"] == interfaceName })?.fields["bridge"] {
            reasons.append(.bridgeMember(bridgeName: bridgeName))
        }
        let matchingAddresses = try await addresses
            .filter { $0.fields["interface"] == interfaceName }
            .compactMap { $0.fields["address"] }
        if !matchingAddresses.isEmpty {
            reasons.append(.hasAddresses(matchingAddresses))
        }
        if try await dhcpClients.contains(where: { $0.fields["interface"] == interfaceName }) {
            reasons.append(.dhcpClient)
        }
        if try await pppoeClients.contains(where: { $0.fields["interface"] == interfaceName }) {
            reasons.append(.pppoeClient)
        }

        return reasons.isEmpty ? nil : PortConflict(interfaceName: interfaceName, reasons: reasons)
    }

    private func finishConnecting(using transport: RouterOSTransport) async {
        activeTransport = transport
        do {
            deviceInfo = try await transport.fetchDeviceInfo()
            interfaces = try await transport.fetchInterfaces()
            // Best-effort: some devices (e.g. CHR/x86 virtual routers) have no physical
            // RouterBOARD at all, so this menu can legitimately be absent — must not fail the
            // whole connection over it.
            routerBoardInfo = try? await Self.fetchRouterBoardInfo(using: transport)
            state = .connected(kind: transport.kind)
        } catch {
            state = .failed(error.localizedDescription)
        }
    }

    /// Field names confirmed live (hEX, RouterOS 6.49.16, see `RouterBoardInfo`'s doc comment).
    private static func fetchRouterBoardInfo(using transport: RouterOSTransport) async throws -> RouterBoardInfo {
        let items = try await transport.fetchMenuItems(menuPath: "/system routerboard", restPath: "system/routerboard")
        guard let fields = items.first?.fields else { throw RouterOSError.invalidResponse("system/routerboard") }
        return RouterBoardInfo(
            model: fields["model"] ?? "unbekannt",
            revision: fields["revision"] ?? "unbekannt",
            serialNumber: fields["serial-number"] ?? "unbekannt",
            firmwareType: fields["firmware-type"] ?? "unbekannt",
            currentFirmware: fields["current-firmware"] ?? "unbekannt",
            minimumFirmware: fields["minimum-firmware"] ?? "unbekannt",
            upgradeFirmware: fields["upgrade-firmware"] ?? "unbekannt"
        )
    }

    func disconnect() async {
        await activeTransport?.disconnect()
        activeTransport = nil
        deviceInfo = nil
        routerBoardInfo = nil
        interfaces = []
        credentials = nil
        hasExpertToolBackedUpThisSession = false
        pendingSSHTrustFingerprint = nil
        state = .idle
    }
}