import XCTest
@testable import MarkBase

final class SamplerTest: XCTestCase {

    var sampler: Sampler!

    override func setUp() {
        super.setUp()
        sampler = Sampler()
    }

    // MARK: - Top-K Sampling

    func testTopKFilter() {
        var logits = [Float](repeating: 0, count: 100)
        logits[10] = 10.0
        logits[20] = 9.0
        logits[30] = 8.0

        let result = sampler.sample(logits: logits, temperature: 1.0, topK: 3, topP: 1.0)
        XCTAssertTrue([10, 20, 30].contains(result),
                      "Should sample from top-3 tokens (\(result))")
    }

    func testTopKAll() {
        var logits = [Float](repeating: 0, count: 100)
        logits[50] = 5.0

        let result = sampler.sample(logits: logits, temperature: 1.0, topK: 100, topP: 1.0)
        XCTAssertGreaterThanOrEqual(result, 0)
        XCTAssertLessThan(result, 100)
    }

    // MARK: - Temperature

    func testTemperatureZeroGreedy() {
        var logits = [Float](repeating: -10, count: 1000)
        logits[42] = 20.0

        let result = sampler.sample(logits: logits, temperature: 0.0, topK: 1, topP: 1.0)
        XCTAssertEqual(result, 42, "Temperature 0 should always select highest logit")
    }

    func testTemperatureHigh() {
        var logits = [Float](repeating: 0, count: 1000)
        logits[0] = 100.0

        let result = sampler.sample(logits: logits, temperature: 10.0, topK: 100, topP: 1.0)
        XCTAssertEqual(result, 0, "Overwhelming logit should still be selected")
    }

    // MARK: - Unused Token Filtering

    func testFilterUnusedTokens() {
        var logits = [Float](repeating: 0, count: 262144)
        logits[258123] = 30.0
        logits[500] = 29.0

        let result = sampler.sample(logits: logits, temperature: 1.0, topK: 50, topP: 0.95, filterUnusedTokens: true)
        XCTAssertLessThan(result, 258000, "Should not sample unused tokens when filtering enabled")
    }

    func testNoUnusedTokenFilter() {
        var logits = [Float](repeating: 0, count: 262144)
        logits[258123] = 50.0
        logits[500] = 0.0

        let result = sampler.sample(logits: logits, temperature: 1.0, topK: 10, topP: 1.0, filterUnusedTokens: false)
        XCTAssertEqual(result, 258123, "Should allow unused tokens when filtering disabled")
    }

    // MARK: - Edge Cases

    func testAllSameLogits() {
        let logits = [Float](repeating: 1.0, count: 1000)
        let result = sampler.sample(logits: logits, temperature: 1.0, topK: 100, topP: 1.0)
        XCTAssertGreaterThanOrEqual(result, 0)
        XCTAssertLessThan(result, 1000)
    }

    func testSingleToken() {
        let logits: [Float] = [5.0]
        let result = sampler.sample(logits: logits, temperature: 1.0, topK: 1, topP: 1.0)
        XCTAssertEqual(result, 0)
    }

    func testExtremeTemperatureZero() {
        var logits = [Float](repeating: -100, count: 100)
        logits[7] = 50.0

        // Greedy: temperature=0, topK=1
        let result = sampler.sample(logits: logits, temperature: 0.0, topK: 1, topP: 1.0)
        XCTAssertEqual(result, 7)
    }
}