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?
    /// True while a lost connection is being silently re-established in the background —
    /// see `startHealthMonitoring()`'s doc comment.
    @Published private(set) var isReconnecting = false
    /// How many REST+SSH reconnect rounds have failed so far this episode — reset to 0
    /// whenever reconnecting starts/succeeds. Purely informational (shown in the banner).
    @Published private(set) var reconnectAttemptCount = 0
    /// Ticks down once a second between reconnect attempts, `nil` while an attempt is
    /// actually in flight or reconnecting isn't happening — lets the banner show "next
    /// attempt in Xs" instead of a bare spinner.
    @Published private(set) var secondsUntilNextReconnectAttempt: Int?
    private var healthMonitorTask: Task<Void, Never>?
    private static let healthCheckInterval: Duration = .seconds(10)
    private static let reconnectRetryInterval: Duration = .seconds(5)

    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
    }

    /// Safety net for when the proactive `verifySSHTrust` check (above) didn't already cover a
    /// dedicated SSH service's host key — confirmed live: a session's first Expert-tab write (its
    /// pre-apply backup, over its own dedicated SSH connection) can still hit
    /// `untrustedSSHHostKey` even after a successful connect, dead-ending as plain error text with
    /// no way to resolve it (same root symptom as Bug 13/37, just at write-time instead of
    /// connect-time). Callers catch `RouterOSError.untrustedSSHHostKey` around their own write
    /// path and call this so the same Verbinden-tab trust dialog appears — trusting there unblocks
    /// a manual retry of whatever write just failed.
    func noteUntrustedSSHHostKey(_ fingerprint: String) {
        pendingSSHTrustFingerprint = fingerprint
    }

    /// 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)
    }

    /// Applies a command over a dedicated, one-shot SSH connection instead of whatever
    /// `activeTransport` currently is — for the small number of menus confirmed live not to work
    /// over RouterOS' REST API at all (not a data/formatting issue on this app's side: the exact
    /// same write succeeds instantly over plain SSH/terminal, REST 500s regardless of which fields
    /// are sent). Confirmed for `/system routerboard mode-button` (2026-09-17); same dedicated-
    /// connection pattern as `BackupService`/`UpdateService`.
    func applyViaSSH(_ command: RouterOSCommand) async throws {
        guard let credentials else { throw RouterOSError.notConnected }
        let transport = SSHTransport(credentials: credentials)
        try await transport.connect()
        do {
            try await transport.apply(command)
            await transport.disconnect()
        } catch {
            await transport.disconnect()
            throw error
        }
    }

    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)
            if healthMonitorTask == nil {
                startHealthMonitoring()
            }
        } catch {
            state = .failed(error.localizedDescription)
        }
    }

    /// Per explicit request: a connection lost mid-session (router rebooted, Wi-Fi
    /// dropped, cable unplugged) should try to recover itself instead of just sitting
    /// there disconnected. Runs for as long as `state` is `.connected`, using a trivial
    /// read (`/system identity`, the smallest possible singleton menu) as a heartbeat —
    /// cheap enough on the router to poll every `healthCheckInterval` indefinitely.
    /// `state` deliberately stays `.connected` throughout a lost-and-recovering episode
    /// so other tabs don't flash to their "not connected" placeholders over what's often
    /// just a brief hiccup; only `isReconnecting` (a small banner) reflects it. Cancelled
    /// in `disconnect()` — an explicit user disconnect must not keep retrying.
    private func startHealthMonitoring() {
        healthMonitorTask = Task { [weak self] in
            while !Task.isCancelled {
                try? await Task.sleep(for: Self.healthCheckInterval)
                guard !Task.isCancelled else { return }
                await self?.checkConnectionHealthAndReconnectIfNeeded()
            }
        }
    }

    private func checkConnectionHealthAndReconnectIfNeeded() async {
        guard case .connected = state, !isReconnecting,
              let activeTransport, let credentials else { return }
        do {
            _ = try await activeTransport.fetchMenuItems(menuPath: "/system identity", restPath: "system/identity")
        } catch {
            await reconnectLoop(credentials: credentials)
        }
    }

    /// Retries REST-then-SSH (same order as a normal `connect()`) every
    /// `reconnectRetryInterval` until one succeeds or the connection is cancelled from
    /// under it (user hit "Trennen" — `disconnect()` cancels `healthMonitorTask`, which
    /// is this loop's own parent Task, and resets `state` to `.idle`, tripping the guard
    /// below on the next iteration regardless).
    private func reconnectLoop(credentials: RouterOSCredentials) async {
        isReconnecting = true
        reconnectAttemptCount = 0
        while !Task.isCancelled {
            guard case .connected = state else { break }
            reconnectAttemptCount += 1
            let rest = RestTransport(credentials: credentials, certificateTrust: certificateTrust)
            if (try? await rest.connect()) != nil {
                await finishConnecting(using: rest)
                resetReconnectState()
                return
            }
            let ssh = SSHTransport(credentials: credentials, hostKeyTrust: sshHostKeyTrust)
            if (try? await ssh.connect()) != nil {
                await finishConnecting(using: ssh)
                resetReconnectState()
                return
            }
            await countdownToNextReconnectAttempt()
        }
        resetReconnectState()
    }

    private func resetReconnectState() {
        isReconnecting = false
        reconnectAttemptCount = 0
        secondsUntilNextReconnectAttempt = nil
    }

    private func countdownToNextReconnectAttempt() async {
        let totalSeconds = Int(Self.reconnectRetryInterval.components.seconds)
        for remaining in stride(from: totalSeconds, through: 1, by: -1) {
            guard !Task.isCancelled else { return }
            secondsUntilNextReconnectAttempt = remaining
            try? await Task.sleep(for: .seconds(1))
        }
        secondsUntilNextReconnectAttempt = nil
    }

    /// 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 {
        healthMonitorTask?.cancel()
        healthMonitorTask = nil
        resetReconnectState()
        await activeTransport?.disconnect()
        activeTransport = nil
        deviceInfo = nil
        routerBoardInfo = nil
        interfaces = []
        credentials = nil
        hasExpertToolBackedUpThisSession = false
        pendingSSHTrustFingerprint = nil
        state = .idle
    }
}