workouts/Workouts Watch App/WorkoutSessionManager.swift

//
// WorkoutSessionManager.swift
// Workouts Watch App
//
// Copyright 2025 Rouslan Zenetl. All Rights Reserved.
//

import Foundation
import HealthKit
import Observation

/// Runs an `HKWorkoutSession` (with a live builder) for the duration of a watch
/// workout.
///
/// When the phone launches us via `startWatchApp(toHandle:)`, watchOS hands the app
/// delegate a workout configuration; starting a session with it grants the
/// freshly-launched app foreground runtime (and keeps it alive while the wrist is
/// down). The attached `HKLiveWorkoutBuilder` makes watchOS collect heart rate and
/// active energy from the sensors, so finishing the session saves a real `HKWorkout`
/// that feeds the Activity rings — and gives us the live HUD values and the metrics
/// we send back to the phone.
@Observable
@MainActor
final class WorkoutSessionManager: NSObject {
    private let healthStore = HKHealthStore()
    private var session: HKWorkoutSession?
    private var builder: HKLiveWorkoutBuilder?

    private(set) var isRunning = false

    /// When the running session started, for the coordinator's session-age rule (a
    /// too-young `.finish` demotes to `.discard`). `distantPast` for a recovered
    /// session — it predates this process, so it's never "young".
    private(set) var sessionStartDate: Date?

    /// Live values for the in-workout HUD. Nil until the first sample arrives (or if
    /// the corresponding read authorization was denied).
    private(set) var currentHeartRate: Double?
    private(set) var currentActiveEnergyKcal: Double?

    /// Seconds accumulated in each of 5 heart-rate zones (low→high). Built on the
    /// watch because it's the only place with the full HR stream. Empty/ignored when
    /// the user's age (hence max HR) is unknown.
    private var hrZoneSeconds = [Double](repeating: 0, count: 5)
    private var lastHRSampleDate: Date?
    private var maxHeartRate: Double?

    /// Guards the terminal paths so a late delegate callback can't double-finish.
    private var isFinishing = false

    /// Prompt for the authorization the live session needs: share workouts + active
    /// energy (so the session can write them), read heart rate + active energy (to
    /// surface them), and read date of birth (to derive HR zones). Called once at
    /// launch so a later phone-initiated launch starts cleanly.
    func requestAuthorization() {
        guard HKHealthStore.isHealthDataAvailable() else { return }
        healthStore.requestAuthorization(
            toShare: WorkoutHealthAuthorization.watchShare,
            read: WorkoutHealthAuthorization.watchRead
        ) { _, _ in }
    }

    /// Start the session for a phone-launched workout. Idempotent — a second launch
    /// while one is already running is ignored.
    func start(with configuration: HKWorkoutConfiguration) {
        guard HKHealthStore.isHealthDataAvailable(), session == nil else { return }
        do {
            let session = try HKWorkoutSession(healthStore: healthStore, configuration: configuration)
            let builder = session.associatedWorkoutBuilder()
            builder.dataSource = HKLiveWorkoutDataSource(healthStore: healthStore, workoutConfiguration: configuration)
            session.delegate = self
            builder.delegate = self
            self.session = session
            self.builder = builder

            resetMetricAccumulators()
            maxHeartRate = computeMaxHeartRate()

            let startDate = Date()
            session.startActivity(with: startDate)
            Task { try? await builder.beginCollection(at: startDate) }
            sessionStartDate = startDate
            isRunning = true
        } catch {
            clear()
        }
    }

    /// Re-adopt a session that outlived a crash or reboot. watchOS keeps the
    /// `HKWorkoutSession` alive across an app death and relaunches us expecting a
    /// recovery call; without it the session is orphaned — its `HKWorkout` never saved
    /// and the live HUD dead. Called once at launch; a no-op when there's nothing to
    /// recover or a fresh start already claimed the slot.
    func recover() {
        guard HKHealthStore.isHealthDataAvailable(), session == nil else { return }
        Task {
            guard let recovered = try? await healthStore.recoverActiveWorkoutSession(),
                  self.session == nil else { return }
            adopt(recovered)
        }
    }

    private func adopt(_ recovered: HKWorkoutSession) {
        let builder = recovered.associatedWorkoutBuilder()
        builder.dataSource = HKLiveWorkoutDataSource(
            healthStore: healthStore, workoutConfiguration: recovered.workoutConfiguration)
        recovered.delegate = self
        builder.delegate = self
        self.session = recovered
        self.builder = builder

        // Zone seconds accumulated before the crash are gone (they lived in memory);
        // resume accumulation from here rather than faking the missing stretch.
        resetMetricAccumulators()
        maxHeartRate = computeMaxHeartRate()

        sessionStartDate = .distantPast
        isRunning = true
    }

    /// Finish the session and save it to Health, returning the captured metrics so the
    /// caller can attach them to the workout and forward to the phone. `totalVolume`
    /// is left nil here — the caller fills it from the workout's logs.
    func finishAndSave() async -> WorkoutMetrics? {
        guard let session, let builder, !isFinishing else { clear(); return nil }
        isFinishing = true
        let endDate = Date()
        session.end()
        do {
            try await builder.endCollection(at: endDate)
            let saved = try await builder.finishWorkout()
            let metrics = makeMetrics(savedWorkout: saved, endDate: endDate)
            clear()
            return metrics
        } catch {
            clear()
            return nil
        }
    }

    /// End the session and throw the data away — nothing is saved to Health. Used when
    /// the workout was discarded rather than completed.
    func discard() {
        guard let session, let builder, !isFinishing else { clear(); return }
        isFinishing = true
        session.end()
        builder.discardWorkout()
        clear()
    }

    // MARK: - Metric helpers

    private func makeMetrics(savedWorkout: HKWorkout?, endDate: Date) -> WorkoutMetrics {
        let hrUnit = HKUnit.count().unitDivided(by: .minute())
        let energy = builder?.statistics(for: HKQuantityType(.activeEnergyBurned))?
            .sumQuantity()?.doubleValue(for: .kilocalorie())
        let hr = builder?.statistics(for: HKQuantityType(.heartRate))
        let zones = (maxHeartRate != nil && hrZoneSeconds.contains { $0 > 0 }) ? hrZoneSeconds : nil

        return WorkoutMetrics(
            activeEnergyKcal: energy,
            avgHeartRate: hr?.averageQuantity()?.doubleValue(for: hrUnit),
            maxHeartRate: hr?.maximumQuantity()?.doubleValue(for: hrUnit),
            minHeartRate: hr?.minimumQuantity()?.doubleValue(for: hrUnit),
            totalVolume: nil,
            hrZoneSeconds: zones,
            healthKitWorkoutUUID: savedWorkout?.uuid.uuidString,
            source: .watch,
            recordedAt: endDate
        )
    }

    /// Pull the latest live stats off the builder for the HUD, and fold the elapsed
    /// time into the right HR zone.
    private func refreshLiveStats() {
        let hrUnit = HKUnit.count().unitDivided(by: .minute())
        let newHR = builder?.statistics(for: HKQuantityType(.heartRate))?
            .mostRecentQuantity()?.doubleValue(for: hrUnit)
        let now = Date()

        if let maxHeartRate, let prevHR = currentHeartRate, let last = lastHRSampleDate {
            let dt = now.timeIntervalSince(last)
            if dt > 0, dt < 60 { hrZoneSeconds[Self.zoneIndex(for: prevHR, maxHR: maxHeartRate)] += dt }
        }
        if let newHR { currentHeartRate = newHR }
        lastHRSampleDate = now

        currentActiveEnergyKcal = builder?.statistics(for: HKQuantityType(.activeEnergyBurned))?
            .sumQuantity()?.doubleValue(for: .kilocalorie())
    }

    /// Bucket an instantaneous heart rate into one of five zones (0–4) by its fraction of the
    /// user's max HR: <60% → 0, 60–70% → 1, 70–80% → 2, 80–90% → 3, ≥90% → 4. Pure (no session
    /// or sensor state), so it's `nonisolated static` and unit-testable off the main actor.
    nonisolated static func zoneIndex(for hr: Double, maxHR: Double) -> Int {
        let ratio = hr / maxHR
        let thresholds = [0.6, 0.7, 0.8, 0.9]
        return thresholds.reduce(0) { $0 + (ratio >= $1 ? 1 : 0) }
    }

    /// Max HR ≈ 220 − age, from the user's date of birth. Nil when unavailable (read
    /// denied or DOB not set), which suppresses zone tracking rather than faking it.
    private func computeMaxHeartRate() -> Double? {
        guard let comps = try? healthStore.dateOfBirthComponents(), let birthYear = comps.year else { return nil }
        let thisYear = Calendar.current.component(.year, from: Date())
        let age = thisYear - birthYear
        guard age > 0, age < 120 else { return nil }
        return Double(220 - age)
    }

    private func resetMetricAccumulators() {
        currentHeartRate = nil
        currentActiveEnergyKcal = nil
        hrZoneSeconds = [Double](repeating: 0, count: 5)
        lastHRSampleDate = nil
        maxHeartRate = nil
    }

    /// The system ended the session out from under us (resource reclaim, watchOS policy).
    /// Salvage what the builder collected into a real Health save instead of dropping it —
    /// the run documents stay active, so no metrics are forwarded to the phone (there is no
    /// completed run to attach them to); the `HKWorkout` itself is what must not be lost.
    private func salvageSystemEndedSession(at endDate: Date) {
        guard let builder, !isFinishing else { clear(); return }
        isFinishing = true
        Task {
            defer { clear() }
            try? await builder.endCollection(at: endDate)
            _ = try? await builder.finishWorkout()
        }
    }

    private func clear() {
        session = nil
        builder = nil
        isRunning = false
        isFinishing = false
        sessionStartDate = nil
        resetMetricAccumulators()
    }
}

// MARK: - HKWorkoutSessionDelegate

extension WorkoutSessionManager: HKWorkoutSessionDelegate {
    nonisolated func workoutSession(
        _ workoutSession: HKWorkoutSession,
        didChangeTo toState: HKWorkoutSessionState,
        from fromState: HKWorkoutSessionState,
        date: Date
    ) {
        guard toState == .ended || toState == .stopped else { return }
        // Only act if we didn't drive the end ourselves (finishAndSave/discard set
        // `isFinishing` and clear on their own). A system-ended session is salvaged
        // into a Health save rather than silently dropped.
        Task { @MainActor in
            if !self.isFinishing { self.salvageSystemEndedSession(at: date) }
        }
    }

    nonisolated func workoutSession(_ workoutSession: HKWorkoutSession, didFailWithError error: Error) {
        Task { @MainActor in self.clear() }
    }
}

// MARK: - HKLiveWorkoutBuilderDelegate

extension WorkoutSessionManager: HKLiveWorkoutBuilderDelegate {
    nonisolated func workoutBuilder(_ workoutBuilder: HKLiveWorkoutBuilder, didCollectDataOf collectedTypes: Set<HKSampleType>) {
        Task { @MainActor in self.refreshLiveStats() }
    }

    nonisolated func workoutBuilderDidCollectEvent(_ workoutBuilder: HKLiveWorkoutBuilder) {}
}