import Foundation

enum AssistantThermalLevel: String, Sendable, Codable, Equatable {
  case nominal
  case fair
  case serious
  case critical
  case unknown
}

struct AssistantResourceSnapshot: Sendable, Equatable {
  let thermalLevel: AssistantThermalLevel
  let lowPowerModeEnabled: Bool
  let availableMemoryBytes: UInt64?

  init(
    thermalLevel: AssistantThermalLevel,
    lowPowerModeEnabled: Bool,
    availableMemoryBytes: UInt64? = nil
  ) {
    self.thermalLevel = thermalLevel
    self.lowPowerModeEnabled = lowPowerModeEnabled
    self.availableMemoryBytes = availableMemoryBytes
  }
}
enum AssistantResourceDecision: Sendable, Equatable {
  case allow
  case warn(String)
  case deny(String)
}

struct ResourceBudgetPolicy: Sendable {
  static func modelLoadDecision(
    for snapshot: AssistantResourceSnapshot
  ) -> AssistantResourceDecision {
    switch snapshot.thermalLevel {
    case .critical:
      return .deny(
        "The iPhone is too hot to load the on-device model safely. Let it cool down and try again."
      )
    case .serious:
      return .warn(
        "The iPhone is running hot. Model loading and generation may be slower until it cools down."
      )
    case .fair, .nominal, .unknown:
      if snapshot.lowPowerModeEnabled {
        return .warn(
          "Low Power Mode is on. On-device generation may be slower."
        )
      }
      return .allow
    }
  }
}