import Atomics

/// A thread-safe, lockless, concurrent sequence which can be appended to and drained.
/// This is not quite a Treiber stack in the sense that single elements cannot
/// be **popped** off of it - Treiber stacks rely on every interaction with a given memory
/// address being a *single atomic read, write or exchange operation*. The Swift atomics
/// package is a bit too anemic for that - we would need to be able to exchange a pointer
/// to the head cell with a pointer to the head cell's child (if any) without reading the
/// head pointer twice.  That's implementable in C or assembly, but not with the access
/// Swift atomics give us.
///
/// So in order to make this work, unlike a traditional treiber stack, we append to the
/// *tail* rather than swapping the head for a new head with the old head as its child
/// Iteration is in FIFO order.  So, Trieber-ish stack, not Treiber stack.
///
/// Implements `Sequence` - the iterator returned may grow in length if the stack is
/// pushed to while iterating, but will never get smaller, even if the stack is emptied.
///
public final class TreiberishStack<Element : Sendable> : Sequence, Sendable {
    // IMPORTANT: If we don't use `ordering: .sequentiallyConsistent` here, we will encounter occasional
    // segfaults when setting the head to nil.
    private let _head : ManagedAtomic<TreiberCell<Element>?>
    private var head : TreiberCell<Element>? { _head.load(ordering: .acquiring) }
    public var first : Element? { head?.element }

    public var last : Element? {
        var h = head
        while h != nil {
            if let n = h!.next {
                h = n
            } else {
                return h!.element
            }
        }
        return nil
    }

    public var count : Int {
        var h = head
        var result = 0
        while let hh = h {
            result += 1
            h = hh.next
        }
        return result
    }
    
    /// isEmpty is constant time
    public var isEmpty : Bool { head == nil }
    
    /// Inherited from Sequence; count returned here will be precisely accurate
    public var underestimatedCount: Int { self.count }

    public subscript(_ index : Int) -> Element? { head?[index] }
    
    public init() { _head  = .init(nil) }
    
    public init<S: Sequence<Element>>(_ sequence : S) where S.Iterator.Element == Element {
        var head : TreiberCell<Element>? = nil
        for element in sequence {
            if let h = head {
                h.push(TreiberCell(element))
            } else {
                head = TreiberCell(element)
            }
        }
        _head = .init(head)
    }
    
    /// Append an element
    public func push(_ element : Element) {
        let cell = TreiberCell(element)
        if let existingHead = _head.compareExchange(expected: nil, desired: cell, ordering: .sequentiallyConsistent).1 {
            existingHead.push(cell)
        }
    }
    
    /// Append many elements in a single atomic operation
    public func push<S: Sequence>(contentsOf : S) where S.Element == Element {
        // Build our tail before we append it
        var append : TreiberCell<Element>? = nil
        for item in contentsOf {
            if let a = append {
                a.push(TreiberCell(item))
            } else {
                append = TreiberCell(item)
            }
        }
        if let a = append {
            if let existingHead = _head.compareExchange(expected: nil, desired: a, ordering: .sequentiallyConsistent).1 {
                existingHead.push(a)
            }
        }
    }
    
    /// Returns an iterator which may grow if the stack is concurrently
    /// appended to, but cannot shrink.  Iteration is in FIFO order.
    public func makeIterator() -> TreiberishStackIterator<Element> {
        .init(cell: head)
    }
    
    public func clear() {
        _head.store(nil, ordering: .releasing)
    }
    
    public func drain() -> [Element] {
        var result : [Element] = []
        if let oldHead = _head.exchange(nil, ordering: .sequentiallyConsistent) {
            // Pending: analyze the compiled code and see if we were really getting
            // tail recursion here, in which case, put it back.
//            oldHead.visit{ result.append($0) }
            var h : TreiberCell<Element>? = oldHead
            while let hh = h {
                result.append(hh.element)
                h = hh.next
            }
        }
        return result
    }
    
    public func exchangeHead(_ newHead : Element) -> [Element] {
        var result : [Element] = []
        if let oldHead = _head.exchange(TreiberCell(newHead), ordering: .sequentiallyConsistent) {
//            oldHead.visit{ result.append($0) }
            var h : TreiberCell<Element>? = oldHead
            while let hh = h {
                result.append(hh.element)
                h = hh.next
            }
        }
        return result
    }
    
    public func map<T>(_ f : (Element) throws -> T) rethrows -> [T] {
        var result : [T] = []
        if let h = head {
            try h.visit { element in
                result.append(try f(element))
            }
        }
        return result
    }

    public func compactMap<T>(_ f : (Element) throws -> T?) rethrows -> [T] {
        var result : [T] = []
        if let h = head {
            try h.visit { element in
                if let t = try f(element) {
                    result.append(t)
                }
            }
        }
        return result
    }

    public func contains(where f : (Element) -> Bool) -> Bool {
        if let h = head {
            var curr : TreiberCell<Element>? = h
            while let c = curr {
                if f(c.element) {
                    return true
                }
                curr = c.next
            }
        }
        return false
    }
}

fileprivate final class TreiberCell<Element : Sendable> : Sendable, AtomicOptionalWrappable, AtomicReference {
    private let _next : ManagedAtomicLazyReference<TreiberCell<Element>> = .init()
    let element : Element
    var next : TreiberCell<Element>? { _next.load() }
    var count : Int { 1 + (next?.count ?? 0) }
    
    var last : Element {
        if let n = next {
            return n.last
        }
        return element
    }
    
    subscript(_ index : Int) -> Element? {
        if index < 0 {
            return nil
        }
        switch index {
        case 0 :
            return element
        default:
            if let n = self.next {
                return n[index - 1]
            }
            return nil
        }
    }
    
    init(_ element: Element) { self.element = element }
    
    func visit(_ f : (Element) throws -> Void) rethrows {
        try f(element)
        if let next = self.next {
            try next.visit(f)
        }
    }
    
    func push(_ cell : TreiberCell<Element>) {
        let next = _next.storeIfNilThenLoad(cell)
        if next !== cell {
            next.push(cell)
        }
    }
}

public struct TreiberishStackIterator<E : Sendable> : IteratorProtocol {
    fileprivate var cell : TreiberCell<E>?
    
    public mutating func next() -> E? {
        guard let current = cell else { return nil }
        defer { cell = current.next }
        return current.element
    }
}