saltylab-ios/SulTee/SulTee/SensorManager.swift

import Foundation
import CoreLocation
import CoreMotion
import MapKit
import Combine

// MARK: - Phone position source

enum PhonePositionSource {
    case uwb(accuracyM: Double)     // robot RTK GPS + UWB offset, computed by Orin
    case gps(accuracyM: Double)     // CoreLocation GPS
    case unknown

    var label: String {
        switch self {
        case .uwb(let a): return String(format: "UWB %.0fcm", a * 100)
        case .gps(let a): return String(format: "GPS %.0fm",  a)
        case .unknown:    return "—"
        }
    }
    var isUWB: Bool { if case .uwb = self { return true }; return false }
}

/// Manages all iPhone sensors, BLE tag streaming, MQTT telemetry, and exposes
/// fused position state for the map and status views.
final class SensorManager: NSObject, ObservableObject {

    // MARK: - Streaming state

    @Published var isStreaming = false
    @Published var wsState: WebSocketClient.ConnectionState = .disconnected

    // MARK: - Sensor rates (Hz)

    @Published var gpsRate:     Double = 0
    @Published var imuRate:     Double = 0
    @Published var headingRate: Double = 0
    @Published var baroRate:    Double = 0

    // MARK: - User (phone) position — fused from UWB-tag or CoreLocation

    /// Best available phone position.  Updated by `updateBestPhonePosition()`.
    @Published var userLocation:       CLLocationCoordinate2D? = nil
    @Published var userBreadcrumbs:    [CLLocationCoordinate2D] = []
    @Published var phonePositionSource: PhonePositionSource = .unknown

    // MARK: - Robot position (saltybot/phone/gps)

    @Published var robotLocation:    CLLocationCoordinate2D? = nil
    @Published var robotBreadcrumbs: [CLLocationCoordinate2D] = []
    @Published var robotSpeed:       Double = 0
    @Published var distanceToRobot:  Double? = nil

    enum DistanceSource {
        case blueUWB    // BLE Range notify from wearable TAG (highest accuracy)
        case mqttUWB    // saltybot/uwb/range from Orin anchors
        case gps        // CoreLocation coordinate diff (fallback)
    }
    @Published var distanceSource: DistanceSource = .gps

    // MARK: - UWB local data (saltybot/uwb/range + saltybot/uwb/position)

    @Published var uwbPosition: UWBPosition? = nil
    @Published var uwbRanges:   [String: UWBRange] = [:]
    @Published var uwbActive:   Bool = false

    // MARK: - Follow settings

    @Published var followMode:   FollowMode   = .gps
    @Published var followPreset: FollowPreset = .medium

    // MARK: - BLE tag

    let ble = BLEManager()

    // MARK: - WebSocket config

    static let defaultOrinURL   = "wss://www.saultee.bot/ws"
    private static let orinURLKey = "orinURL"

    private(set) var ws: WebSocketClient

    var orinURLString: String {
        get { UserDefaults.standard.string(forKey: Self.orinURLKey) ?? Self.defaultOrinURL }
        set { UserDefaults.standard.set(newValue, forKey: Self.orinURLKey) }
    }

    // MARK: - MQTT

    private let mqtt = MQTTClient(config: .init(
        host:     "192.168.87.29",
        port:     1883,
        username: "mqtt_seb",
        password: "mqtt_pass",
        clientID: "saul-t-mote-\(UUID().uuidString.prefix(8))"
    ))

    private static let iosGPSTopic        = "saltybot/ios/gps"
    private static let robotGPSTopic      = "saltybot/phone/gps"
    private static let uwbRangeTopic      = "saltybot/uwb/range"
    private static let uwbPositionTopic   = "saltybot/uwb/position"
    private static let uwbTagPosTopic     = "saltybot/uwb/tag/position"   // Orin-fused phone position
    private static let followModeTopic    = "saltybot/follow/mode"
    private static let followRangeTopic   = "saltybot/follow/range"
    private static let maxBreadcrumbs     = 60
    private static let uwbStaleSeconds    = 3.0

    // MARK: - Internal sensor state

    private var lastKnownLocation:   CLLocation?
    private var lastKnownMotion:     CMDeviceMotion?

    /// Orin-fused phone absolute position (RTK GPS + UWB offset).
    private var uwbTagPosition:      (coord: CLLocationCoordinate2D, accuracyM: Double, ts: Date)?

    // MARK: - Timers

    private var mqttGPSTimer:       Timer?   // 1 Hz MQTT publish
    private var bleGPSTimer:        Timer?   // 5 Hz BLE GPS write
    private var bleIMUTimer:        Timer?   // 10 Hz BLE IMU write
    private var uwbStalenessTimer:  Timer?
    private var rateTimer:          Timer?

    // MARK: - CoreMotion / CoreLocation

    private let locationManager = CLLocationManager()
    private let motionManager   = CMMotionManager()
    private let altimeter       = CMAltimeter()
    private var cancellables    = Set<AnyCancellable>()
    private var bleCancellables = Set<AnyCancellable>()

    // MARK: - Rate counters

    private var gpsCounts:     [Date] = []
    private var imuCounts:     [Date] = []
    private var headingCounts: [Date] = []
    private var baroCounts:    [Date] = []

    // MARK: - Init

    override init() {
        // Migrate: if the stored URL is the old Tailscale IP, replace with the new WSS endpoint.
        if let saved = UserDefaults.standard.string(forKey: Self.orinURLKey),
           saved.contains("100.64.0.2") {
            UserDefaults.standard.removeObject(forKey: Self.orinURLKey)
        }
        let urlStr = UserDefaults.standard.string(forKey: Self.orinURLKey) ?? Self.defaultOrinURL
        self.ws = WebSocketClient(url: URL(string: urlStr) ?? URL(string: Self.defaultOrinURL)!)
        super.init()

        locationManager.delegate = self
        locationManager.desiredAccuracy = kCLLocationAccuracyBestForNavigation
        locationManager.distanceFilter = kCLDistanceFilterNone
        locationManager.allowsBackgroundLocationUpdates = true
        locationManager.pausesLocationUpdatesAutomatically = false
        locationManager.showsBackgroundLocationIndicator = true

        ws.$state
            .receive(on: DispatchQueue.main)
            .assign(to: \.wsState, on: self)
            .store(in: &cancellables)

        mqtt.onMessage = { [weak self] topic, payload in
            guard let self else { return }
            switch topic {
            case Self.robotGPSTopic:    self.handleRobotGPS(payload)
            case Self.uwbRangeTopic:    self.handleUWBRange(payload)
            case Self.uwbPositionTopic: self.handleUWBPosition(payload)
            case Self.uwbTagPosTopic:   self.handleUWBTagPosition(payload)
            default: break
            }
        }

        // Start BLE timers whenever the tag connects; stop when it disconnects
        ble.$connectionState
            .receive(on: DispatchQueue.main)
            .sink { [weak self] state in
                guard let self else { return }
                if state == .connected {
                    self.startBLETimers()
                    // Ensure sensors are running (needed even without Follow-Me mode)
                    self.ensureSensorsRunning()
                } else {
                    self.stopBLETimers()
                    // Re-evaluate distance source now that BLE anchors are gone
                    self.updateDistanceToRobot()
                }
            }
            .store(in: &bleCancellables)

        // Re-evaluate distance on every Range notify from the wearable TAG
        ble.$anchors
            .receive(on: DispatchQueue.main)
            .sink { [weak self] _ in self?.updateDistanceToRobot() }
            .store(in: &bleCancellables)
    }

    // MARK: - Public control (Follow-Me / WebSocket)

    func startStreaming() {
        guard !isStreaming else { return }
        isStreaming = true
        ws.connect()
        mqtt.connect()
        mqtt.subscribe(topic: Self.robotGPSTopic)
        mqtt.subscribe(topic: Self.uwbRangeTopic)
        mqtt.subscribe(topic: Self.uwbPositionTopic)
        mqtt.subscribe(topic: Self.uwbTagPosTopic)
        ensureSensorsRunning()
        startRateTimer()
        startMQTTGPSTimer()
        startUWBStalenessTimer()
        publishFollowMode()
        publishFollowPreset()
    }

    func stopStreaming() {
        guard isStreaming else { return }
        isStreaming = false
        ws.disconnect()
        mqtt.disconnect()
        // Keep sensors running if BLE is connected; otherwise stop
        if !ble.isConnected { stopSensors() }
        rateTimer?.invalidate();    rateTimer = nil
        mqttGPSTimer?.invalidate(); mqttGPSTimer = nil
        uwbStalenessTimer?.invalidate(); uwbStalenessTimer = nil
    }

    func updateURL(_ urlString: String) {
        guard !isStreaming else { return }
        orinURLString = urlString
        if let url = URL(string: urlString), url.scheme?.hasPrefix("ws") == true {
            ws.url = url
        }
    }

    func setFollowMode(_ mode: FollowMode) {
        followMode = mode; publishFollowMode()
    }

    func setFollowPreset(_ preset: FollowPreset) {
        followPreset = preset; publishFollowPreset()
    }

    // MARK: - Sensor lifecycle

    private func ensureSensorsRunning() {
        locationManager.requestAlwaysAuthorization()
        locationManager.startUpdatingLocation()
        locationManager.startUpdatingHeading()
        if !motionManager.isDeviceMotionActive { startIMU() }
        if CMAltimeter.isRelativeAltitudeAvailable() { startBarometer() }
    }

    private func stopSensors() {
        locationManager.stopUpdatingLocation()
        locationManager.stopUpdatingHeading()
        motionManager.stopDeviceMotionUpdates()
        altimeter.stopRelativeAltitudeUpdates()
    }

    // MARK: - BLE streaming timers

    private func startBLETimers() {
        stopBLETimers()
        // GPS → tag at 5 Hz (200 ms)
        bleGPSTimer = Timer.scheduledTimer(withTimeInterval: 0.2, repeats: true) { [weak self] _ in
            guard let self, let loc = self.lastKnownLocation else { return }
            self.ble.sendGPS(BLEPackets.gpsPacket(from: loc))
        }
        // IMU → tag at 10 Hz (100 ms)
        bleIMUTimer = Timer.scheduledTimer(withTimeInterval: 0.1, repeats: true) { [weak self] _ in
            guard let self, let motion = self.lastKnownMotion else { return }
            self.ble.sendIMU(BLEPackets.imuPacket(from: motion))
        }
    }

    private func stopBLETimers() {
        bleGPSTimer?.invalidate(); bleGPSTimer = nil
        bleIMUTimer?.invalidate(); bleIMUTimer = nil
    }

    // MARK: - MQTT GPS publish (1 Hz)

    private func startMQTTGPSTimer() {
        mqttGPSTimer = Timer.scheduledTimer(withTimeInterval: 1.0, repeats: true) { [weak self] _ in
            self?.publishGPSToMQTT()
        }
    }

    private func publishGPSToMQTT() {
        guard let loc = lastKnownLocation else { return }
        let payload: [String: Any] = [
            "ts":          loc.timestamp.timeIntervalSince1970,
            "lat":         loc.coordinate.latitude,
            "lon":         loc.coordinate.longitude,
            "alt_m":       loc.altitude,
            "accuracy_m":  max(0, loc.horizontalAccuracy),
            "speed_ms":    max(0, loc.speed),
            "bearing_deg": loc.course >= 0 ? loc.course : 0.0,
            "provider":    "gps"
        ]
        guard let data = try? JSONSerialization.data(withJSONObject: payload),
              let json = String(data: data, encoding: .utf8) else { return }
        mqtt.publish(topic: Self.iosGPSTopic, payload: json)
    }

    // MARK: - MQTT follow publish helpers

    private func publishFollowMode()   { mqtt.publish(topic: Self.followModeTopic,  payload: followMode.mqttPayload) }
    private func publishFollowPreset() { mqtt.publish(topic: Self.followRangeTopic, payload: followPreset.mqttPayload) }

    // MARK: - Incoming MQTT handlers

    private func handleRobotGPS(_ payload: String) {
        guard let data = payload.data(using: .utf8),
              let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let lat  = json["lat"] as? Double,
              let lon  = json["lon"] as? Double else { return }
        let coord = CLLocationCoordinate2D(latitude: lat, longitude: lon)
        robotLocation = coord
        robotSpeed    = (json["speed_ms"] as? Double) ?? 0
        appendBreadcrumb(coord, to: &robotBreadcrumbs)
        updateDistanceToRobot()
    }

    private func handleUWBRange(_ payload: String) {
        guard let data     = payload.data(using: .utf8),
              let json     = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let anchorID = json["anchor_id"] as? String,
              let rangeM   = json["range_m"] as? Double else { return }
        uwbRanges[anchorID] = UWBRange(anchorID: anchorID, rangeMetres: rangeM, timestamp: Date())
        updateDistanceToRobot()
    }

    private func handleUWBPosition(_ payload: String) {
        guard let data = payload.data(using: .utf8),
              let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let x   = json["x"] as? Double,
              let y   = json["y"] as? Double,
              let z   = json["z"] as? Double else { return }
        uwbPosition = UWBPosition(x: x, y: y, z: z, timestamp: Date())
        uwbActive   = true
    }

    /// Orin-fused phone absolute position: robot RTK GPS + UWB offset.
    /// This is the most accurate phone position when UWB is in range.
    private func handleUWBTagPosition(_ payload: String) {
        guard let data      = payload.data(using: .utf8),
              let json      = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let lat       = json["lat"] as? Double,
              let lon       = json["lon"] as? Double else { return }
        let accuracy = (json["accuracy_m"] as? Double) ?? 0.02   // default 2 cm for UWB
        let coord = CLLocationCoordinate2D(latitude: lat, longitude: lon)
        uwbTagPosition = (coord: coord, accuracyM: accuracy, ts: Date())
        updateBestPhonePosition()
    }

    // MARK: - Phone position source selection

    /// Selects the best available phone position:
    ///   1. UWB-derived (saltybot/uwb/tag/position) if fresh < 3 s
    ///   2. CoreLocation GPS fallback
    private func updateBestPhonePosition() {
        if let uwb = uwbTagPosition,
           Date().timeIntervalSince(uwb.ts) < Self.uwbStaleSeconds {
            // Robot RTK + UWB offset is the authority
            let coord = uwb.coord
            if userLocation != coord {
                userLocation = coord
                appendBreadcrumb(coord, to: &userBreadcrumbs)
                updateDistanceToRobot()
            }
            phonePositionSource = .uwb(accuracyM: uwb.accuracyM)
        } else if let loc = lastKnownLocation {
            let coord = loc.coordinate
            if userLocation != coord {
                userLocation = coord
                appendBreadcrumb(coord, to: &userBreadcrumbs)
                updateDistanceToRobot()
            }
            phonePositionSource = .gps(accuracyM: max(0, loc.horizontalAccuracy))
        }
    }

    // MARK: - UWB staleness watchdog

    private func startUWBStalenessTimer() {
        uwbStalenessTimer = Timer.scheduledTimer(withTimeInterval: 1.0, repeats: true) { [weak self] _ in
            guard let self else { return }
            let cutoff = Date().addingTimeInterval(-Self.uwbStaleSeconds)
            if let pos = self.uwbPosition, pos.timestamp < cutoff { self.uwbActive = false }
            self.uwbRanges = self.uwbRanges.filter { $0.value.timestamp > cutoff }
            // Re-evaluate phone position source when UWB tag position may have gone stale
            self.updateBestPhonePosition()
        }
    }

    // MARK: - Breadcrumbs + distance

    private func appendBreadcrumb(_ coord: CLLocationCoordinate2D,
                                  to list: inout [CLLocationCoordinate2D]) {
        list.append(coord)
        if list.count > Self.maxBreadcrumbs { list.removeFirst() }
    }

    /// Update distanceToRobot using the highest-accuracy source available:
    ///   1. BLE Range notify (wearable TAG) — cm-accurate, freshness checked via AnchorInfo.isStale
    ///   2. MQTT UWB ranges (saltybot/uwb/range from Orin anchors)
    ///   3. GPS coordinate diff — fallback only
    func updateDistanceToRobot() {
        // Priority 1: fresh BLE anchor data from the wearable TAG
        let freshAnchors = ble.anchors.filter { !$0.isStale }
        if !freshAnchors.isEmpty {
            let closest = freshAnchors.min(by: { $0.rangeMetres < $1.rangeMetres })!
            distanceToRobot = closest.rangeMetres
            distanceSource  = .blueUWB
            return
        }

        // Priority 2: MQTT UWB ranges from Orin anchors
        let freshRanges = uwbRanges.values.filter {
            Date().timeIntervalSince($0.timestamp) < Self.uwbStaleSeconds
        }
        if let closest = freshRanges.min(by: { $0.rangeMetres < $1.rangeMetres }) {
            distanceToRobot = closest.rangeMetres
            distanceSource  = .mqttUWB
            return
        }

        // Priority 3: GPS coordinate diff
        guard let user = userLocation, let robot = robotLocation else { return }
        let a = CLLocation(latitude: user.latitude,  longitude: user.longitude)
        let b = CLLocation(latitude: robot.latitude, longitude: robot.longitude)
        distanceToRobot = a.distance(from: b)
        distanceSource  = .gps
    }

    // MARK: - IMU / Barometer

    private func startIMU() {
        guard motionManager.isDeviceMotionAvailable else { return }
        motionManager.deviceMotionUpdateInterval = 1.0 / 100.0
        motionManager.startDeviceMotionUpdates(to: .main) { [weak self] motion, _ in
            guard let self, let motion else { return }
            self.lastKnownMotion = motion
            self.recordEvent(in: &self.imuCounts)
            self.ws.send([
                "type": "imu",
                "timestamp": Date().timeIntervalSince1970,
                "data": [
                    "accel":    ["x": motion.userAcceleration.x, "y": motion.userAcceleration.y, "z": motion.userAcceleration.z],
                    "gyro":     ["x": motion.rotationRate.x,     "y": motion.rotationRate.y,     "z": motion.rotationRate.z],
                    "attitude": ["roll": motion.attitude.roll,    "pitch": motion.attitude.pitch, "yaw": motion.attitude.yaw],
                    "gravity":  ["x": motion.gravity.x,           "y": motion.gravity.y,          "z": motion.gravity.z],
                    "magneticField": [
                        "x": motion.magneticField.field.x,
                        "y": motion.magneticField.field.y,
                        "z": motion.magneticField.field.z,
                        "accuracy": motion.magneticField.accuracy.rawValue
                    ]
                ]
            ])
        }
    }

    private func startBarometer() {
        guard CMAltimeter.isRelativeAltitudeAvailable() else { return }
        altimeter.startRelativeAltitudeUpdates(to: .main) { [weak self] data, _ in
            guard let self, let data else { return }
            self.recordEvent(in: &self.baroCounts)
            self.ws.send([
                "type": "baro",
                "timestamp": Date().timeIntervalSince1970,
                "data": ["relativeAltitude": data.relativeAltitude.doubleValue,
                         "pressure":         data.pressure.doubleValue]
            ])
        }
    }

    // MARK: - Rate tracking

    private func startRateTimer() {
        rateTimer = Timer.scheduledTimer(withTimeInterval: 1.0, repeats: true) { [weak self] _ in
            self?.updateRates()
        }
    }

    private func recordEvent(in list: inout [Date]) { list.append(Date()) }

    private func updateRates() {
        let cutoff = Date().addingTimeInterval(-1.0)
        gpsCounts     = gpsCounts.filter     { $0 > cutoff }
        imuCounts     = imuCounts.filter     { $0 > cutoff }
        headingCounts = headingCounts.filter { $0 > cutoff }
        baroCounts    = baroCounts.filter    { $0 > cutoff }
        DispatchQueue.main.async {
            self.gpsRate     = Double(self.gpsCounts.count)
            self.imuRate     = Double(self.imuCounts.count)
            self.headingRate = Double(self.headingCounts.count)
            self.baroRate    = Double(self.baroCounts.count)
        }
    }
}

// MARK: - CLLocationManagerDelegate

extension SensorManager: CLLocationManagerDelegate {

    func locationManager(_ manager: CLLocationManager, didUpdateLocations locations: [CLLocation]) {
        guard let loc = locations.last else { return }
        lastKnownLocation = loc
        recordEvent(in: &gpsCounts)
        // Let the source-selection logic decide whether to use this or UWB-derived position
        updateBestPhonePosition()

        ws.send([
            "type": "gps",
            "timestamp": loc.timestamp.timeIntervalSince1970,
            "data": [
                "latitude":           loc.coordinate.latitude,
                "longitude":          loc.coordinate.longitude,
                "altitude":           loc.altitude,
                "horizontalAccuracy": loc.horizontalAccuracy,
                "verticalAccuracy":   loc.verticalAccuracy,
                "speed":              loc.speed,
                "speedAccuracy":      loc.speedAccuracy,
                "course":             loc.course,
                "courseAccuracy":     loc.courseAccuracy
            ]
        ])
    }

    func locationManager(_ manager: CLLocationManager, didUpdateHeading newHeading: CLHeading) {
        recordEvent(in: &headingCounts)
        ws.send([
            "type": "heading",
            "timestamp": Date().timeIntervalSince1970,
            "data": [
                "magneticHeading": newHeading.magneticHeading,
                "trueHeading":     newHeading.trueHeading,
                "headingAccuracy": newHeading.headingAccuracy,
                "x": newHeading.x, "y": newHeading.y, "z": newHeading.z
            ]
        ])
    }

    func locationManager(_ manager: CLLocationManager, didFailWithError error: Error) {
        print("[Location] error: \(error)")
    }

    func locationManagerDidChangeAuthorization(_ manager: CLLocationManager) {
        switch manager.authorizationStatus {
        case .authorizedAlways, .authorizedWhenInUse:
            if isStreaming || ble.isConnected {
                manager.startUpdatingLocation()
                manager.startUpdatingHeading()
            }
        default: break
        }
    }
}