← poker-tutor

import Foundation

/// The strength of a five-card poker hand, packed into one comparable number:
/// category in the high bits, then five rank nibbles as tiebreakers.
public struct HandRank: Comparable, Hashable, Sendable, CustomStringConvertible {
    public let value: UInt32

    public enum Category: UInt32, Comparable, Sendable, CaseIterable {
        case highCard = 0, pair, twoPair, threeOfAKind, straight, flush, fullHouse, fourOfAKind, straightFlush

        public static func < (a: Category, b: Category) -> Bool { a.rawValue < b.rawValue }

        public var name: String {
            switch self {
            case .highCard: return "high card"
            case .pair: return "a pair"
            case .twoPair: return "two pair"
            case .threeOfAKind: return "three of a kind"
            case .straight: return "a straight"
            case .flush: return "a flush"
            case .fullHouse: return "a full house"
            case .fourOfAKind: return "four of a kind"
            case .straightFlush: return "a straight flush"
            }
        }
    }

    public init(value: UInt32) { self.value = value }

    init(_ category: Category, _ k1: UInt32 = 0, _ k2: UInt32 = 0, _ k3: UInt32 = 0, _ k4: UInt32 = 0, _ k5: UInt32 = 0) {
        value = (category.rawValue << 20) | (k1 << 16) | (k2 << 12) | (k3 << 8) | (k4 << 4) | k5
    }

    public var category: Category { Category(rawValue: value >> 20)! }

    /// The five tiebreaker ranks, most significant first. Zeros mean "unused".
    public var kickers: [UInt32] {
        [(value >> 16) & 0xF, (value >> 12) & 0xF, (value >> 8) & 0xF, (value >> 4) & 0xF, value & 0xF]
    }

    public static func < (a: HandRank, b: HandRank) -> Bool { a.value < b.value }

    public var description: String {
        let ranks = kickers.filter { $0 != 0 }.compactMap { Rank(rawValue: UInt8($0))?.letter }.map(String.init)
        return ranks.isEmpty ? category.name : "\(category.name) (\(ranks.joined(separator: "-")))"
    }
}

public enum HandEvaluator {

    /// Best five-card hand out of five, six or seven cards.
    public static func evaluate(_ cards: [Card]) -> HandRank {
        precondition(cards.count >= 5 && cards.count <= 7, "evaluate expects 5-7 cards")

        var rankCount = [UInt8](repeating: 0, count: 15)   // indexed by Rank.rawValue
        var suitCount = [UInt8](repeating: 0, count: 4)
        var suitMask = [UInt16](repeating: 0, count: 4)
        var rankMask: UInt16 = 0

        for card in cards {
            let r = card.rank
            let s = Int(card.suit.rawValue)
            rankCount[Int(r.rawValue)] += 1
            suitCount[s] += 1
            suitMask[s] |= r.bit
            rankMask |= r.bit
        }

        var best = HandRank(value: 0)

        // Flushes and straight flushes.
        if let flushSuit = (0..<4).first(where: { suitCount[$0] >= 5 }) {
            let mask = suitMask[flushSuit]
            if let high = straightHigh(mask) {
                return HandRank(.straightFlush, UInt32(high))   // nothing beats this
            }
            let top = topRanks(mask, 5)
            best = HandRank(.flush, top[0], top[1], top[2], top[3], top[4])
        }

        // Rank-count hands.
        var quad: UInt32 = 0
        var trips: [UInt32] = []
        var pairs: [UInt32] = []
        for r in stride(from: 14, through: 2, by: -1) {
            switch rankCount[r] {
            case 4: quad = UInt32(r)
            case 3: trips.append(UInt32(r))
            case 2: pairs.append(UInt32(r))
            default: break
            }
        }

        if quad != 0 {
            let kicker = topRanks(rankMask & ~Rank(rawValue: UInt8(quad))!.bit, 1)[0]
            return max(best, HandRank(.fourOfAKind, quad, kicker))
        }

        if let trip = trips.first {
            // A second set plays as the pair; otherwise the best actual pair does.
            let pairRank = trips.count > 1 ? trips[1] : pairs.first
            if let pairRank {
                return max(best, HandRank(.fullHouse, trip, pairRank))
            }
        }

        if best.category == .flush { return best }

        if let high = straightHigh(rankMask) {
            return HandRank(.straight, UInt32(high))
        }

        if let trip = trips.first {
            let rest = topRanks(rankMask & ~Rank(rawValue: UInt8(trip))!.bit, 2)
            return HandRank(.threeOfAKind, trip, rest[0], rest[1])
        }

        if pairs.count >= 2 {
            let used = Rank(rawValue: UInt8(pairs[0]))!.bit | Rank(rawValue: UInt8(pairs[1]))!.bit
            let kicker = topRanks(rankMask & ~used, 1)[0]
            return HandRank(.twoPair, pairs[0], pairs[1], kicker)
        }

        if let pair = pairs.first {
            let rest = topRanks(rankMask & ~Rank(rawValue: UInt8(pair))!.bit, 3)
            return HandRank(.pair, pair, rest[0], rest[1], rest[2])
        }

        let top = topRanks(rankMask, 5)
        return HandRank(.highCard, top[0], top[1], top[2], top[3], top[4])
    }

    /// Highest card of a five-in-a-row inside a 13-bit rank mask, or nil. Handles the wheel (5-4-3-2-A).
    public static func straightHigh(_ mask: UInt16) -> UInt8? {
        var high: UInt8 = 14
        while high >= 6 {
            let window: UInt16 = 0b11111 << UInt16(high - 6)
            if mask & window == window { return high }
            high -= 1
        }
        let wheel: UInt16 = 0b1_0000_0000_1111   // A,5,4,3,2
        return mask & wheel == wheel ? 5 : nil
    }

    /// The `count` highest ranks present in a mask, as rank raw values.
    public static func topRanks(_ mask: UInt16, _ count: Int) -> [UInt32] {
        var out: [UInt32] = []
        out.reserveCapacity(count)
        var r = 14
        while r >= 2 && out.count < count {
            if mask & Rank(rawValue: UInt8(r))!.bit != 0 { out.append(UInt32(r)) }
            r -= 1
        }
        while out.count < count { out.append(0) }
        return out
    }
}