import Foundation

enum PuzzleGenerationError: Error, LocalizedError {
    case unsupportedSize(Int)
    case failed(size: Int, attempts: Int)

    var errorDescription: String? {
        switch self {
        case .unsupportedSize(let size):
            return "Grid size \(size) is not supported."
        case .failed(let size, let attempts):
            return "Puzzle generation failed for \(size)x\(size) after \(attempts) attempts."
        }
    }
}

struct SeededRandomNumberGenerator: RandomNumberGenerator {
    private var state: UInt64

    init(seed: UInt64) {
        self.state = seed == 0 ? 0x9E3779B97F4A7C15 : seed
    }

    mutating func next() -> UInt64 {
        state &+= 0x9E3779B97F4A7C15
        var z = state
        z = (z ^ (z >> 30)) &* 0xBF58476D1CE4E5B9
        z = (z ^ (z >> 27)) &* 0x94D049BB133111EB
        return z ^ (z >> 31)
    }

    mutating func randomDouble() -> Double {
        let value = next() >> 11
        return Double(value) / Double(1 << 53)
    }

    mutating func weightedLowIndex(count: Int) -> Int {
        guard count > 1 else {
            return 0
        }

        // The Python version sampled a beta distribution to favor smaller rectangles.
        // This power curve keeps the same bias without pulling in a statistics library.
        let fraction = pow(randomDouble(), 2.2)
        return min(Int(fraction * Double(count)), count - 1)
    }
}

enum ShikakuGenerator {
    static let maximumAttempts = 200
    private static let maximumRandomSeed: UInt64 = 2_147_483_648

    static func generate(size: Int, seed requestedSeed: UInt64? = nil) throws -> ShikakuPuzzle {
        guard size > 1 else {
            throw PuzzleGenerationError.unsupportedSize(size)
        }

        let seed = requestedSeed ?? UInt64.random(in: 0...maximumRandomSeed)
        var rng = SeededRandomNumberGenerator(seed: seed)

        for _ in 0..<maximumAttempts {
            if let puzzle = attempt(size: size, seed: seed, rng: &rng) {
                return puzzle
            }
        }

        throw PuzzleGenerationError.failed(size: size, attempts: maximumAttempts)
    }

    private static func attempt(size: Int, seed: UInt64, rng: inout SeededRandomNumberGenerator) -> ShikakuPuzzle? {
        var grid = Array(repeating: Array(repeating: -1, count: size), count: size)
        var rects: [SolutionRect] = []
        var uncovered = (0..<size).flatMap { row in
            (0..<size).map { col in GridPoint(row: row, col: col) }
        }

        uncovered.shuffle(using: &rng)

        while !uncovered.isEmpty {
            let start = uncovered[0]

            if grid[start.row][start.col] != -1 {
                uncovered.removeFirst()
                continue
            }

            var candidates: [SolutionRect] = []

            for height in 1...(size - start.row) {
                for width in 1...(size - start.col) {
                    if height == 1 && width == 1 {
                        continue
                    }

                    for rowOffset in 0..<height {
                        for colOffset in 0..<width {
                            let row0 = start.row - rowOffset
                            let col0 = start.col - colOffset
                            let row1 = row0 + height - 1
                            let col1 = col0 + width - 1

                            guard row0 >= 0, col0 >= 0, row1 < size, col1 < size else {
                                continue
                            }

                            if rectangleIsUncovered(row0: row0, col0: col0, row1: row1, col1: col1, grid: grid) {
                                candidates.append(SolutionRect(row: row0, col: col0, height: height, width: width))
                            }
                        }
                    }
                }
            }

            if candidates.isEmpty {
                guard absorbIsolatedCell(start, size: size, grid: &grid, rects: &rects, uncovered: &uncovered, rng: &rng) else {
                    return nil
                }
                continue
            }

            let keyedCandidates = candidates.map { rect in
                (rect: rect, tieBreak: rng.randomDouble())
            }
            let sortedCandidates = keyedCandidates
                .sorted { lhs, rhs in
                    if lhs.rect.area == rhs.rect.area {
                        return lhs.tieBreak < rhs.tieBreak
                    }
                    return lhs.rect.area < rhs.rect.area
                }
                .map(\.rect)

            let chosen = sortedCandidates[rng.weightedLowIndex(count: sortedCandidates.count)]
            let rectID = rects.count

            for cell in chosen.cells {
                grid[cell.row][cell.col] = rectID
            }

            rects.append(chosen)
            uncovered.removeAll { grid[$0.row][$0.col] != -1 }
        }

        guard (0..<size).allSatisfy({ row in (0..<size).allSatisfy { grid[row][$0] != -1 } }) else {
            return nil
        }

        var clues: [GridPoint: Int] = [:]
        for rect in rects {
            guard let clueCell = rect.cells.randomElement(using: &rng) else {
                return nil
            }
            clues[clueCell] = rect.area
        }

        return ShikakuPuzzle(size: size, clues: clues, solution: rects, seed: seed)
    }

    private static func rectangleIsUncovered(
        row0: Int,
        col0: Int,
        row1: Int,
        col1: Int,
        grid: [[Int]]
    ) -> Bool {
        for row in row0...row1 {
            for col in col0...col1 where grid[row][col] != -1 {
                return false
            }
        }
        return true
    }

    private static func absorbIsolatedCell(
        _ cell: GridPoint,
        size: Int,
        grid: inout [[Int]],
        rects: inout [SolutionRect],
        uncovered: inout [GridPoint],
        rng: inout SeededRandomNumberGenerator
    ) -> Bool {
        var neighbors = [
            GridPoint(row: cell.row - 1, col: cell.col),
            GridPoint(row: cell.row + 1, col: cell.col),
            GridPoint(row: cell.row, col: cell.col - 1),
            GridPoint(row: cell.row, col: cell.col + 1)
        ]
        neighbors.shuffle(using: &rng)

        for neighbor in neighbors {
            guard (0..<size).contains(neighbor.row),
                  (0..<size).contains(neighbor.col),
                  grid[neighbor.row][neighbor.col] != -1 else {
                continue
            }

            let rectID = grid[neighbor.row][neighbor.col]
            let old = rects[rectID]
            let oldBounds = old.bounds
            let newRow0 = min(oldBounds.startRow, cell.row)
            let newCol0 = min(oldBounds.startCol, cell.col)
            let newRow1 = max(oldBounds.endRow, cell.row)
            let newCol1 = max(oldBounds.endCol, cell.col)

            let expandedCells = CellBounds(
                startRow: newRow0,
                startCol: newCol0,
                endRow: newRow1,
                endCol: newCol1
            ).cells

            guard expandedCells.allSatisfy({ grid[$0.row][$0.col] == -1 || grid[$0.row][$0.col] == rectID }) else {
                continue
            }

            for expandedCell in expandedCells {
                grid[expandedCell.row][expandedCell.col] = rectID
            }

            rects[rectID] = SolutionRect(
                row: newRow0,
                col: newCol0,
                height: newRow1 - newRow0 + 1,
                width: newCol1 - newCol0 + 1
            )
            uncovered.removeAll { grid[$0.row][$0.col] != -1 }
            return true
        }

        return false
    }
}