Tusker/Packages/TTTKit/Sources/TTTKit/AI/GameModel.swift

//
//  GameModel.swift
//  TTTKit
//
//  Created by Shadowfacts on 12/21/22.
//

import Foundation
import GameKit

public class GameModel: NSObject, NSCopying, GKGameModel {
    private var controller: GameController
    
    public init(controller: GameController) {
        self.controller = controller
    }
    
    // MARK: GKGameModel
    
    public var players: [GKGameModelPlayer]? {
        [Player.x, Player.o]
    }
    
    public var activePlayer: GKGameModelPlayer? {
        switch controller.state {
        case .playAnywhere(let mark), .playSpecific(let mark, column: _, row: _):
            switch mark {
            case .x:
                return Player.x
            case .o:
                return Player.o
            }
        case .end(_):
            return nil
        }
    }
    
    public func setGameModel(_ gameModel: GKGameModel) {
        let other = (gameModel as! GameModel).controller
        self.controller = GameController(state: other.state, board: other.board)
    }
    
    public func gameModelUpdates(for player: GKGameModelPlayer) -> [GKGameModelUpdate]? {
        guard let player = player as? Player else {
            return nil
        }
        let mark = player.mark
        
        switch controller.state {
        case .playAnywhere:
            return updatesForPlayAnywhere(mark: mark)
        case .playSpecific(_, column: let col, row: let row):
            return updatesForPlaySpecific(mark: mark, board: (col, row))
        case .end:
            return nil
        }
    }
    
    public func apply(_ gameModelUpdate: GKGameModelUpdate) {
        let update = gameModelUpdate as! Update
        switch controller.state {
        case .playAnywhere(update.mark), .playSpecific(update.mark, column: update.subBoard.column, row: update.subBoard.row):
            break
        default:
            return
        }
        controller.play(on: update.subBoard, column: update.column, row: update.row)
    }
    
    public func score(for player: GKGameModelPlayer) -> Int {
        guard let player = player as? Player else {
            return .min
        }
        
        var score = 0
        for column in 0..<3 {
            for row in 0..<3 {
                let subBoard = controller.board.getSubBoard(column: column, row: row)
                if let win = subBoard.win {
                    if win.mark == player.mark {
                        score += 10
                    } else {
                        score -= 5
                    }
                } else {
                    score += 4 * subBoard.potentialWinCount(for: player.mark)
                    score -= 2 * subBoard.potentialWinCount(for: player.mark.next)
                }
            }
        }
        if case .playAnywhere(let mark) = controller.state {
            if mark == player.mark {
                score += 5
            }
        }
        score += 8 * controller.board.potentialWinCount(for: player.mark)
        score -= 4 * controller.board.potentialWinCount(for: player.mark.next)
        return score
    }
    
    public func isWin(for player: GKGameModelPlayer) -> Bool {
        let mark = (player as! Player).mark
        return controller.board.win?.mark == mark
    }
    
    private func updatesForPlayAnywhere(mark: Mark) -> [Update] {
        var updates = [Update]()
        for boardColumn in 0..<3 {
            for boardRow in 0..<3 {
                let subBoard = controller.board.getSubBoard(column: boardColumn, row: boardRow)
                guard !subBoard.ended else { continue }
                
                for column in 0..<3 {
                    for row in 0..<3 {
                        guard subBoard[column, row] == nil else { continue }
                        updates.append(Update(mark: mark, subBoard: (boardColumn, boardRow), column: column, row: row))
                    }
                }
            }
        }
        return updates
    }
    
    private func updatesForPlaySpecific(mark: Mark, board: (column: Int, row: Int)) -> [Update] {
        let subBoard = controller.board.getSubBoard(column: board.column, row: board.row)
        var updates = [Update]()
        for column in 0..<3 {
            for row in 0..<3 {
                guard subBoard[column, row] == nil else { continue }
                updates.append(Update(mark: mark, subBoard: board, column: column, row: row))
            }
        }
        return updates
    }
    
    // MARK: NSCopying
    
    public func copy(with zone: NSZone? = nil) -> Any {
        return GameModel(controller: GameController(state: controller.state, board: controller.board))
    }
    
}

extension Board {
    func potentialWinCount(for mark: Mark) -> Int {
        return Win.allPoints.filter { points in
            let empty = points.filter { self[$0] == nil }.count
            let matching = points.filter { self[$0] == mark }.count
            return matching == 2 && empty == 1
        }.count
    }
}