RouterOS/RouterOSAssistant/Features/Wizard/Steps/Connect/ConnectViewModel.swift

import Foundation

@MainActor
final class ConnectViewModel: ObservableObject {
    @Published var host: String = "192.168.88.1"
    @Published var username: String = "admin"
    @Published var password: String = ""
    @Published var rememberPassword: Bool = true
    @Published private(set) var savedRouters: [SavedRouter] = []

    @Published private(set) var packageUpdateInfo: PackageUpdateInfo?
    @Published private(set) var isCheckingForUpdates = false
    @Published var updateCheckError: String?

    @Published private(set) var isInstallingUpdate = false
    @Published var installUpdateError: String?
    @Published private(set) var didSendUpdateInstall = false

    @Published private(set) var isUpgradingFirmware = false
    @Published var firmwareUpgradeError: String?
    @Published var firmwareUpgradeResultMessage: String?

    @Published private(set) var isRebooting = false
    @Published var rebootError: String?
    @Published private(set) var didSendReboot = false

    @Published private(set) var interfaceTraffic: [String: InterfaceTraffic] = [:]

    let connectionService: ConnectionService
    private let keychain = KeychainService()
    private let updateService = UpdateService()
    private let savedRoutersStore = SavedRoutersStore()
    private let trafficMonitor = InterfaceTrafficMonitor()
    private var trafficPollingTask: Task<Void, Never>?

    init(connectionService: ConnectionService) {
        self.connectionService = connectionService
    }

    func onAppear() {
        savedRouters = savedRoutersStore.load()
        if let saved = keychain.loadPassword(forHost: host, username: username) {
            password = saved
        }
    }

    /// Fills in a known router's host/username (and its remembered password, if any) without
    /// connecting yet — the user still reviews/confirms via the normal "Verbinden" button.
    func selectSavedRouter(_ router: SavedRouter) {
        host = router.host
        username = router.username
        password = keychain.loadPassword(forHost: router.host, username: router.username) ?? ""
    }

    func renameSavedRouter(_ id: SavedRouter.ID, to newName: String) {
        savedRouters = savedRoutersStore.rename(id, to: newName)
    }

    func updateSavedRouterLocation(_ id: SavedRouter.ID, to newLocation: String) {
        savedRouters = savedRoutersStore.updateLocation(id, to: newLocation)
    }

    func removeSavedRouter(_ id: SavedRouter.ID) {
        savedRouters = savedRoutersStore.remove(id)
    }

    func connect() {
        let credentials = RouterOSCredentials(host: host, username: username, password: password)
        if rememberPassword {
            keychain.save(password: password, forHost: host, username: username)
        }
        Task {
            await connectionService.connect(with: credentials)
            if case .connected = connectionService.state {
                let defaultName = connectionService.deviceInfo?.boardName ?? host
                savedRouters = savedRoutersStore.recordSuccessfulConnection(
                    host: host, username: username, defaultName: defaultName
                )
            }
        }
    }

    /// Starts (or restarts, if already running) periodic traffic polling for whatever interfaces
    /// `connectionService.interfaces` currently lists — re-read every tick so newly-appearing
    /// interfaces (e.g. a VLAN added via the Setup wizard) get picked up without restarting the
    /// poll. 3s cadence: fast enough to feel "live" for a status dot, not so fast it noticeably
    /// loads a home router's CPU with a constant stream of SSH round-trips.
    func startTrafficPolling(credentials: RouterOSCredentials) {
        stopTrafficPolling()
        trafficPollingTask = Task {
            while !Task.isCancelled {
                let names = connectionService.interfaces.map(\.name)
                if !names.isEmpty {
                    let traffic = await trafficMonitor.fetchTraffic(interfaceNames: names, for: credentials)
                    if !Task.isCancelled {
                        interfaceTraffic = traffic
                    }
                }
                try? await Task.sleep(for: .seconds(3))
            }
        }
    }

    func stopTrafficPolling() {
        trafficPollingTask?.cancel()
        trafficPollingTask = nil
        interfaceTraffic = [:]
        Task { await trafficMonitor.disconnect() }
    }

    func trustAndRetry(fingerprint: String) {
        Task {
            await connectionService.trustCurrentCertificateAndRetry(fingerprint: fingerprint)
        }
    }

    func trustSSHHostKeyAndRetry(fingerprint: String) {
        Task {
            await connectionService.trustCurrentSSHHostKeyAndRetry(fingerprint: fingerprint)
        }
    }

    func checkForUpdates(for credentials: RouterOSCredentials) {
        isCheckingForUpdates = true
        updateCheckError = nil
        Task {
            do {
                packageUpdateInfo = try await updateService.checkForUpdates(for: credentials)
            } catch {
                updateCheckError = error.localizedDescription
            }
            isCheckingForUpdates = false
        }
    }

    func installUpdate(for credentials: RouterOSCredentials) {
        isInstallingUpdate = true
        installUpdateError = nil
        didSendUpdateInstall = false
        Task {
            do {
                try await updateService.installUpdate(for: credentials)
                didSendUpdateInstall = true
            } catch {
                installUpdateError = error.localizedDescription
            }
            isInstallingUpdate = false
        }
    }

    func upgradeRouterboardFirmware(for credentials: RouterOSCredentials) {
        isUpgradingFirmware = true
        firmwareUpgradeError = nil
        firmwareUpgradeResultMessage = nil
        Task {
            do {
                firmwareUpgradeResultMessage = try await updateService.upgradeRouterboardFirmware(for: credentials)
            } catch {
                firmwareUpgradeError = error.localizedDescription
            }
            isUpgradingFirmware = false
        }
    }

    func reboot(for credentials: RouterOSCredentials) {
        isRebooting = true
        rebootError = nil
        didSendReboot = false
        Task {
            do {
                try await updateService.reboot(for: credentials)
                didSendReboot = true
            } catch {
                rebootError = error.localizedDescription
            }
            isRebooting = false
        }
    }
}