import Foundation
import Observation
import RastrilloNative

/// Sends one account's outbox. The app supplies `send`, which turns a batch
/// of one lane into requests; the drainer decides when to call it, and what
/// an error it throws means, so every app reacts the same way:
///
/// | Thrown | Reaction |
/// |---|---|
/// | `.offline`, path down | stop; wait for a connectivity change, with no timer |
/// | `.offline`, path up (a timeout, a lost response) | back off, retry the same batch |
/// | `.unauthorized` | session ended: stop every lane; never hold or drop |
/// | `.rateLimited(after)` | back off at least `after` |
/// | `.server`, `.transport`, `.invalidResponse` | back off, retry the same batch |
/// | `BlobStoreError.unavailable` | when thrown by send, retry this lane; hold after three consecutive throws |
/// | `.forbidden`, `.notFound`, `.conflict`, `.tooLarge`, `.rejected` | hold the batch rather than spin |
///
/// A failure to acknowledge or hold after a send is a `.retry`: the server
/// may already have the batch, and each app's server-side idempotency makes
/// sending it again harmless.
///
/// At most one batch is in flight, across all lanes: the simplest correct
/// ordering against one server.
@MainActor @Observable public final class OutboxDrainer {
  public enum Disposition: Sendable, Equatable {
    case sent
    case retry
    case hold(Hold)
  }

  /// Where the drainer is. `SyncStatus` reads it.
  public enum Phase: Sendable, Equatable {
    case idle
    case sending
    case waitingForNetwork
    case backingOff
    case sessionEnded
    case storeUnavailable
  }

  /// What one `drain()` ended on.
  public enum Outcome: Sendable, Equatable {
    /// Nothing left that can be sent.
    case idle
    case waitingForNetwork
    /// Retry after this delay (a connectivity change may cut it short).
    case backingOff(TimeInterval)
    /// The server’s Retry-After: neither connectivity nor "Try again" crosses it.
    case rateLimited(TimeInterval)
    case sessionEnded
    case storeUnavailable
  }

  /// Returns a disposition per entry. An entry it leaves out was not sent
  /// this time (for example, a merge that had to fall back to the first
  /// entry): it goes back to sendable with no backoff.
  public typealias Send = @MainActor (_ lane: String, _ entries: [OutboxEntry], _ transport: NativeTransport) async throws -> [UUID: Disposition]

  public let outbox: Outbox
  public private(set) var phase: Phase = .idle
  public private(set) var summary = OutboxSummary()
  public private(set) var sessionEnded = false
  private let transport: @MainActor () -> NativeTransport?
  private let connectivity: NativeConnectivity
  private let send: Send
  private let maxBatchBytes: Int?
  private var backoff: Backoff
  @ObservationIgnored private var knownStatuses: [String: SyncStatus] = [:]
  private var consecutiveThrows: [String: Int] = [:]
  private var draining = false
  private var drainWaiters: [CheckedContinuation<Outcome, Never>] = []
  private var worker: Task<Void, Never>?
  private var observation: Task<Void, Never>?
  private var networkObservation: Task<Void, Never>?
  private var timer: Task<Void, Never>?
  private var wake: AsyncStream<Void>.Continuation?
  private var retryDeadline: TimeInterval?
  private var serverDeadline: TimeInterval?
  private var now: TimeInterval { ProcessInfo.processInfo.systemUptime }

  public init(outbox: Outbox, transport: @escaping @MainActor () -> NativeTransport?, connectivity: NativeConnectivity,
              maxBatchBytes: Int? = nil, backoff: Backoff = Backoff(), send: @escaping Send) {
    self.outbox = outbox; self.transport = transport; self.connectivity = connectivity
    self.maxBatchBytes = maxBatchBytes; self.backoff = backoff; self.send = send
  }
  public func start() {
    guard worker == nil else { return }
    let (events, continuation) = AsyncStream<Void>.makeStream(bufferingPolicy: .bufferingNewest(1))
    wake = continuation
    observation = Task { [weak self, outbox] in
      for await summary in await outbox.changes {
        guard let self else { return }
        if self.summary != summary { self.summary = summary; self.kick() }
      }
    }
    let changes = connectivity.changes
    networkObservation = Task { [weak self] in
      for await _ in changes {
        guard let self else { return }
        self.backoff.reset(); self.retryDeadline = nil; self.kick()
      }
    }
    worker = Task { [weak self] in
      for await _ in events {
        guard let self, !Task.isCancelled else { return }
        if !self.connectivity.isPathSatisfied {
          self.timer?.cancel(); self.timer = nil
          self.phase = .waitingForNetwork; continue
        }
        let deadline = max(self.retryDeadline ?? 0, self.serverDeadline ?? 0)
        if deadline > self.now { self.schedule(deadline - self.now); continue }
        switch await self.drain() {
        case .backingOff(let delay): self.retryDeadline = self.now + delay; self.schedule(delay)
        case .rateLimited(let delay): self.schedule(delay)
        default: break
        }
      }
    }
    kick()
  }
  private func schedule(_ delay: TimeInterval) {
    timer?.cancel()
    timer = Task { [weak self] in
      do { try await Task.sleep(for: .seconds(delay)) } catch { return }
      self?.kick()
    }
  }
  public func stop() {
    worker?.cancel(); observation?.cancel(); networkObservation?.cancel(); timer?.cancel(); wake?.finish()
    worker = nil; observation = nil; networkObservation = nil; timer = nil; wake = nil
  }
  public func kick() { wake?.yield() }
  public func retryNow() { retryDeadline = nil; backoff.reset(); kick() }
  public func sessionRestored() { sessionEnded = false; retryNow() }
  private func retry() -> Outcome { phase = .backingOff; return .backingOff(backoff.next()) }

  @discardableResult public func drain() async -> Outcome {
    if draining { return await withCheckedContinuation { drainWaiters.append($0) } }
    draining = true
    let outcome = await drainOnce()
    draining = false
    let waiters = drainWaiters; drainWaiters.removeAll()
    for waiter in waiters { waiter.resume(returning: outcome) }
    return outcome
  }
  private func drainOnce() async -> Outcome {
    do { try await outbox.load() }
    catch { summary = await outbox.summary; phase = .storeUnavailable; return .storeUnavailable }
    summary = await outbox.summary
    if sessionEnded { phase = .sessionEnded; return .sessionEnded }
    if !connectivity.isPathSatisfied { phase = .waitingForNetwork; return .waitingForNetwork }
    if let deadline = serverDeadline, deadline > now { phase = .backingOff; return .rateLimited(deadline - now) }
    serverDeadline = nil
    var retryLanes: Set<String> = []
    while !Task.isCancelled {
      summary = await outbox.summary
      var batch: [OutboxEntry] = []
      do {
        for lane in summary.lanes.keys.sorted() where !retryLanes.contains(lane) {
          batch = try await outbox.beginSend(lane: lane, maxBytes: maxBatchBytes)
          if !batch.isEmpty { break }
        }
      } catch { phase = .storeUnavailable; return .storeUnavailable }
      guard let first = batch.first else {
        if !retryLanes.isEmpty { return retry() }
        phase = .idle; return .idle
      }
      let ids = batch.map(\.id)
      guard let transport = transport() else {
        await outbox.endSend(ids, failed: false); sessionEnded = true; phase = .sessionEnded; return .sessionEnded
      }
      phase = .sending
      do {
        let dispositions: [UUID: Disposition]
        do {
          dispositions = try await send(first.lane, batch, transport)
          consecutiveThrows[first.lane] = nil
        }
        catch let error as NativeClientError {
          consecutiveThrows[first.lane] = nil
          let hold: Hold?
          switch error {
          case .forbidden, .notFound: hold = Hold(.accessLost)
          case .conflict(let data): hold = Hold(.conflict(data))
          case .tooLarge: hold = Hold(.tooLarge)
          case .rejected(let status, _): hold = Hold(.rejected(String(status)))
          default: hold = nil
          }
          guard let hold else { throw error }
          dispositions = Dictionary(uniqueKeysWithValues: ids.map { ($0, .hold(hold)) })
        }
        catch {
          guard !Task.isCancelled else { throw error }
          let count = (consecutiveThrows[first.lane] ?? 0) + 1
          consecutiveThrows[first.lane] = count
          // An app decoder/programming failure belongs to this lane, not the
          // account. Three consecutive throws require an explicit resolution.
          let disposition: Disposition = count >= 3 ? .hold(Hold(.rejected("send"))) : .retry
          dispositions = Dictionary(uniqueKeysWithValues: ids.map { ($0, disposition) })
          if count >= 3 { consecutiveThrows[first.lane] = nil }
        }
        var mustRetry = false; var progressed = false
        for entry in batch {
          switch dispositions[entry.id] {
          case .sent: try await outbox.ack([entry.id]); progressed = true
          case .hold(let hold): try await outbox.hold([entry.id], hold); progressed = true
          case .retry: await outbox.endSend([entry.id]); mustRetry = true
          case nil: break
          }
        }
        await outbox.endSend(ids, failed: false)
        summary = await outbox.summary
        if mustRetry || !progressed { retryLanes.insert(first.lane); continue }
        backoff.reset(); retryDeadline = nil
      } catch {
        await outbox.endSend(ids); summary = await outbox.summary
        switch error {
        case NativeClientError.unauthorized:
          sessionEnded = true; phase = .sessionEnded; return .sessionEnded
        case NativeClientError.offline where !connectivity.isPathSatisfied:
          phase = .waitingForNetwork; return .waitingForNetwork
        case NativeClientError.rateLimited(let after):
          let delay = backoff.next(retryAfter: after); serverDeadline = now + delay
          phase = .backingOff; return .rateLimited(delay)
        default: return retry()
        }
      }
      if !connectivity.isPathSatisfied { phase = .waitingForNetwork; return .waitingForNetwork }
    }
    phase = .idle; return .idle
  }
  public func status(lane: String? = nil, localFailure: SyncFailure? = nil,
                     count: (OutboxSummary) -> Int = { $0.pendingLanes }) -> SyncStatus {
    let state = lane.map { summary.only($0) } ?? summary
    let offline = transport().map { connectivity.isOffline($0.origin) } ?? !connectivity.isPathSatisfied
    let key = lane.map { "lane:" + $0 } ?? "account"
    if state.unavailable, localFailure == nil, let known = knownStatuses[key] { return known }
    let status = SyncStatus(summary: state, sending: state.inFlight > 0, sessionEnded: sessionEnded,
                            offline: offline, backingOff: phase == .backingOff, localFailure: localFailure, count: count)
    // A locked store retains the last observable status. An app's explicit
    // enqueue failure still reports input that has not reached durable storage.
    if !state.unavailable, localFailure == nil { knownStatuses[key] = status }
    return status
  }
}