SwimBuddy/targets/watch/WorkoutManager.swift

import Foundation
import SwiftUI
import HealthKit
import WatchKit
import WatchConnectivity

class WorkoutManager: ObservableObject {
    @Published var isWorkoutActive = false
    @Published var savedWorkouts: [StoredWorkout] = []
    @Published var completedWorkouts: [StoredWorkout] = []
    @Published var selectedWorkout: StoredWorkout?
    @Published var currentSetIndex = 0
    @Published var currentRep = 0
    @Published var currentLapDistance = 0
    @Published var elapsedSeconds: TimeInterval = 0
    @Published var heartRate: Double = 0
    @Published var isResting = false
    @Published var restTimeRemaining: TimeInterval = 0

    private var workoutTimer: Timer?
    private var restTimer: Timer?
    private var healthStore: HKHealthStore?
    private var activeWorkout: HKWorkoutSession?
    private var builder: HKWorkoutBuilder?
    private var lapStartTime: Date?
    private var completedLaps: [CompletedLap] = []

    init() {
        healthStore = HKHealthStore()
    }

    // MARK: - Data Persistence

    func loadSavedWorkouts() {
        let defaults = UserDefaults.standard

        if let data = defaults.data(forKey: "savedWorkouts"),
           let workouts = try? JSONDecoder().decode([StoredWorkout].self, from: data) {
            savedWorkouts = workouts
        }

        if let data = defaults.data(forKey: "completedWorkouts"),
           let workouts = try? JSONDecoder().decode([StoredWorkout].self, from: data) {
            completedWorkouts = workouts
        }
    }

    func saveWorkoutsToStorage() {
        let defaults = UserDefaults.standard
        if let data = try? JSONEncoder().encode(savedWorkouts) {
            defaults.set(data, forKey: "savedWorkouts")
        }
        if let data = try? JSONEncoder().encode(completedWorkouts) {
            defaults.set(data, forKey: "completedWorkouts")
        }
    }

    func addWorkoutFromPhone(_ plan: StoredWorkout) {
        savedWorkouts.append(plan)
        saveWorkoutsToStorage()
    }

    // MARK: - Workout Selection

    func selectWorkout(_ workout: StoredWorkout) {
        selectedWorkout = workout
    }

    // MARK: - Workout Control

    func startWorkout() {
        guard selectedWorkout != nil else { return }

        isWorkoutActive = true
        currentSetIndex = 0
        currentRep = 0
        elapsedSeconds = 0
        completedLaps = []

        startHealthKitWorkout()
        startTimer()
        WKInterfaceDevice.current().play(.start)
    }

    func pauseWorkout() {
        workoutTimer?.invalidate()
        WKInterfaceDevice.current().play(.notification)
    }

    func resumeWorkout() {
        startTimer()
        WKInterfaceDevice.current().play(.start)
    }

    func endWorkout() {
        workoutTimer?.invalidate()
        restTimer?.invalidate()
        isWorkoutActive = false

        endHealthKitWorkout()
        syncCompletedWorkout()
        WKInterfaceDevice.current().play(.success)
    }

    // MARK: - Lap Tracking

    func completeLap() {
        guard let workout = selectedWorkout else { return }

        let now = Date()
        let duration = lapStartTime.map { now.timeIntervalSince($0) } ?? 0

        let lap = CompletedLap(
            setPosition: currentSetIndex + 1,
            repNumber: currentRep + 1,
            distance: workout.sets[currentSetIndex].distancePerRep,
            duration: duration,
            heartRate: Int(heartRate)
        )
        completedLaps.append(lap)

        WKInterfaceDevice.current().play(.click)

        advanceToNextRep()
    }

    private func advanceToNextRep() {
        guard let workout = selectedWorkout else { return }
        let currentSet = workout.sets[currentSetIndex]

        if currentRep + 1 < currentSet.repCount {
            currentRep += 1
            lapStartTime = Date()
        } else if currentSetIndex + 1 < workout.sets.count {
            startRestPeriod()
        } else {
            endWorkout()
        }
    }

    // MARK: - Rest Period

    private func startRestPeriod() {
        guard let workout = selectedWorkout else { return }
        let currentSet = workout.sets[currentSetIndex]

        guard currentSet.restType != "none", currentSet.restDuration > 0 else {
            moveToNextSet()
            return
        }

        isResting = true
        restTimeRemaining = currentSet.restDuration

        WKInterfaceDevice.current().play(.notification)

        restTimer = Timer.scheduledTimer(withTimeInterval: 1, repeats: true) { [weak self] _ in
            guard let self = self else { return }
            self.restTimeRemaining -= 1

            if self.restTimeRemaining <= 0 {
                self.restTimer?.invalidate()
                self.moveToNextSet()
            }
        }
    }

    private func moveToNextSet() {
        isResting = false
        currentSetIndex += 1
        currentRep = 0
        lapStartTime = Date()

        if currentSetIndex >= (selectedWorkout?.sets.count ?? 0) {
            endWorkout()
        } else {
            WKInterfaceDevice.current().play(.start)
        }
    }

    // MARK: - Timer

    private func startTimer() {
        lapStartTime = Date()
        workoutTimer = Timer.scheduledTimer(withTimeInterval: 1, repeats: true) { [weak self] _ in
            self?.elapsedSeconds += 1
        }
    }

    // MARK: - HealthKit

    private func startHealthKitWorkout() {
        guard let healthStore = healthStore else { return }

        let config = HKWorkoutConfiguration()
        config.activityType = .swimming
        config.locationType = .indoor

        do {
            activeWorkout = try HKWorkoutSession(healthStore: healthStore, configuration: config)
            let workoutBuilder = HKWorkoutBuilder(healthStore: healthStore, configuration: config, device: nil)
            self.builder = workoutBuilder
            activeWorkout?.startActivity(with: Date())
            workoutBuilder.beginCollection(withStart: Date()) { _, _ in }
        } catch {
            print("Failed to start HealthKit workout: \(error)")
        }
    }

    private func endHealthKitWorkout() {
        activeWorkout?.end()
        builder?.endCollection(withEnd: Date()) { [weak self] _, _ in
            self?.builder?.finishWorkout { _, _ in }
        }
    }

    // MARK: - Sync to Phone

    private func syncCompletedWorkout() {
        let workoutSummary: [String: Any] = [
            "type": "workoutCompleted",
            "workoutId": selectedWorkout?.id ?? "",
            "laps": completedLaps.map { lap in
                [
                    "setPosition": lap.setPosition,
                    "repNumber": lap.repNumber,
                    "distance": lap.distance,
                    "duration": lap.duration,
                    "heartRate": lap.heartRate ?? 0,
                ] as [String: Any]
            },
            "duration": elapsedSeconds,
        ]

        var completed = completedWorkouts
        var summary = selectedWorkout
        summary?.completedAt = ISO8601DateFormatter().string(from: Date())
        if let s = summary {
            completed.insert(s, at: 0)
        }
        if let data = try? JSONEncoder().encode(completed) {
            UserDefaults.standard.set(data, forKey: "completedWorkouts")
        }

        if WCSession.default.isReachable {
            WCSession.default.sendMessage(workoutSummary, replyHandler: nil) { error in
                print("Failed to sync to phone: \(error)")
            }
        } else {
            WCSession.default.transferUserInfo(workoutSummary)
        }
    }
}

// MARK: - Models

struct StoredWorkout: Codable, Identifiable {
    let id: String
    let name: String?
    let poolLength: Double
    let sets: [StoredSet]
    var completedAt: String?
}

struct StoredSet: Codable {
    let stroke: String
    let distancePerRep: Int
    let repCount: Int
    let restType: String
    let restDuration: TimeInterval
}

struct CompletedLap {
    let setPosition: Int
    let repNumber: Int
    let distance: Int
    let duration: TimeInterval
    let heartRate: Int?
}