//
//  AsyncOneShot.swift
//  Tun
//
//  Created by Codex on 2026/4/28.
//

import Foundation

enum AsyncOneShotError: Error {
    case timeout
}

actor AsyncOneShot<Value: Sendable> {
    private struct Waiter {
        let continuation: CheckedContinuation<Value, Error>
        let timeoutTask: Task<Void, Never>?
    }

    private var result: Result<Value, Error>?
    private var waiters: [UUID: Waiter] = [:]

    func wait() async throws -> Value {
        try await self.waitInternal(timeout: nil, timeoutError: AsyncOneShotError.timeout)
    }

    func wait(timeout: Duration, timeoutError: Error = AsyncOneShotError.timeout) async throws -> Value {
        try await self.waitInternal(timeout: timeout, timeoutError: timeoutError)
    }

    func succeed(_ value: Value) {
        self.complete(.success(value))
    }

    func fail(_ error: Error) {
        self.complete(.failure(error))
    }

    private func waitInternal(timeout: Duration?, timeoutError: Error) async throws -> Value {
        let id = UUID()

        return try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { continuation in
                self.addWaiter(
                    id: id,
                    continuation: continuation,
                    timeout: timeout,
                    timeoutError: timeoutError
                )
            }
        } onCancel: {
            Task {
                await self.cancelWaiter(id: id, throwing: CancellationError())
            }
        }
    }

    private func addWaiter(
        id: UUID,
        continuation: CheckedContinuation<Value, Error>,
        timeout: Duration?,
        timeoutError: Error
    ) {
        if let result {
            continuation.resume(with: result)
            return
        }

        let timeoutTask = timeout.map { timeout in
            Task {
                do {
                    try await Task.sleep(for: timeout)
                    if !Task.isCancelled {
                        self.cancelWaiter(id: id, throwing: timeoutError)
                    }
                } catch {
                    return
                }
            }
        }

        self.waiters[id] = Waiter(continuation: continuation, timeoutTask: timeoutTask)
    }

    private func complete(_ result: Result<Value, Error>) {
        guard self.result == nil else {
            return
        }

        self.result = result

        let waiters = self.waiters
        self.waiters.removeAll()

        for waiter in waiters.values {
            waiter.timeoutTask?.cancel()
            waiter.continuation.resume(with: result)
        }
    }

    private func cancelWaiter(id: UUID, throwing error: Error) {
        guard let waiter = self.waiters.removeValue(forKey: id) else {
            return
        }

        waiter.timeoutTask?.cancel()
        waiter.continuation.resume(throwing: error)
    }
}