-- RLDiseaseTransitionTests.lua
-- Purpose: Tests for Disease state machine transitions
-- Author: Ritter
-- Validates: Disease.new() constructor, treatment cycle, immunity expiry,
--            natural recovery, fatality curve
-- Pattern: Mock disease types with known values, direct state manipulation

RLDiseaseTransitionTests = {}

local Log = RmLogging.getLogger("RLRM")

--- Run all disease transition tests.
--- Tests the Disease:onPeriodChanged() state machine using mock disease types.
--- @param assert function Assert closure from RLTestRunner
function RLDiseaseTransitionTests:runAll(assert)
    Log:info("--- Disease Transition Tests ---")

    -- Save state
    local savedDiseasesEnabled = g_diseaseManager.diseasesEnabled
    g_diseaseManager.diseasesEnabled = true

    self:testConstructor(assert)
    self:testTreatmentCycle(assert)
    self:testImmunityExpiry(assert)
    self:testNaturalRecovery(assert)
    self:testFatalityCurve(assert)

    -- Restore state
    g_diseaseManager.diseasesEnabled = savedDiseasesEnabled

    -- Cleanup
    RLTestUtils:cleanupTestAnimals()
end

-- =========================================================================
-- Constructor
-- =========================================================================

--- Verify Disease.new() sets initial state correctly.
--- @param assert function Assert closure from RLTestRunner
function RLDiseaseTransitionTests:testConstructor(assert)
    Log:info("--- Disease Constructor ---")

    local dt = RLTestUtils:createMockDiseaseType()

    -- Default: not carrier, genes=0
    local d1 = Disease.new(dt, false, 0)
    assert("Disease.new default: cured=false", d1.cured == false, false, d1.cured)
    assert("Disease.new default: time=-1", d1.time == -1, -1, d1.time)
    assert("Disease.new default: immunity=0", d1.immunity == 0, 0, d1.immunity)

    -- Carrier with genes=1
    local d2 = Disease.new(dt, true, 1)
    assert("Disease.new carrier: isCarrier=true", d2.isCarrier == true, true, d2.isCarrier)
    assert("Disease.new carrier: genes=1", d2.genes == 1, 1, d2.genes)

    -- Affected with genes=2
    local d3 = Disease.new(dt, false, 2)
    assert("Disease.new affected: genes=2", d3.genes == 2, 2, d3.genes)
end

-- =========================================================================
-- Treatment Cycle
-- =========================================================================

--- Test treatment → cured → immunity path.
--- Mock treatment: cost=100, duration=2.
--- @param assert function Assert closure from RLTestRunner
function RLDiseaseTransitionTests:testTreatmentCycle(assert)
    Log:info("--- Disease Treatment Cycle ---")

    local dt = RLTestUtils:createMockDiseaseType({
        treatment = { cost = 100, duration = 2 },
        immunity = 6,
        recovery = false,  -- suppress natural recovery for isolation
    })

    local disease = Disease.new(dt, false, 0)
    disease.beingTreated = true
    disease.treatmentDuration = 2

    -- Create animal to receive onPeriodChanged calls
    local animal = RLTestUtils:createTestAnimal("COW")
    table.insert(animal.diseases, disease)

    -- First period: treatmentDuration 2→1, not cured yet
    local died1, cost1 = disease:onPeriodChanged(animal, false)
    assert("treatment period 1: treatmentDuration decrements",
        disease.treatmentDuration == 1, 1, disease.treatmentDuration)
    assert("treatment period 1: cost returned",
        cost1 == 100, 100, cost1)

    -- Second period: treatmentDuration 1→0, cured
    local died2, cost2 = disease:onPeriodChanged(animal, false)
    assert("treatment period 2: cured=true", disease.cured == true, true, disease.cured)
    assert("treatment period 2: immunity set to type.immunity",
        disease.immunity == 6, 6, disease.immunity)

    RLTestUtils:cleanupTestAnimals()
end

-- =========================================================================
-- Immunity Expiry
-- =========================================================================

--- Test immunity countdown and disease removal on expiry.
--- @param assert function Assert closure from RLTestRunner
function RLDiseaseTransitionTests:testImmunityExpiry(assert)
    Log:info("--- Disease Immunity Expiry ---")

    local dt = RLTestUtils:createMockDiseaseType({ recovery = false })

    local disease = Disease.new(dt, false, 0)
    disease.cured = true
    disease.immunity = 1

    -- Create animal with this disease in its array
    local animal = RLTestUtils:createTestAnimal("COW")
    table.insert(animal.diseases, disease)
    local countBefore = #animal.diseases

    -- Call onPeriodChanged → immunity=0 → disease removed
    disease:onPeriodChanged(animal, false)
    assert("immunity expiry: immunity reaches 0",
        disease.immunity == 0, 0, disease.immunity)
    assert("immunity expiry: disease removed from animal",
        #animal.diseases == countBefore - 1, countBefore - 1, #animal.diseases)

    RLTestUtils:cleanupTestAnimals()
end

-- =========================================================================
-- Natural Recovery
-- =========================================================================

--- Test natural recovery counter and statistical cure.
--- Recovery increments each period. When recovery >= type.recovery AND
--- math.random() >= 0.25, disease is cured (75% chance per eligible period).
--- @param assert function Assert closure from RLTestRunner
function RLDiseaseTransitionTests:testNaturalRecovery(assert)
    Log:info("--- Disease Natural Recovery ---")

    local dt = RLTestUtils:createMockDiseaseType({
        recovery = 3,
        fatality = { {time = 100, value = 0} },  -- no fatality
    })

    -- Test recovery counter increment (deterministic)
    local d1 = Disease.new(dt, false, 0)
    local animal1 = RLTestUtils:createTestAnimal("COW")
    table.insert(animal1.diseases, d1)

    d1:onPeriodChanged(animal1, false)  -- time -1→0, recovery 0→1
    d1:onPeriodChanged(animal1, false)  -- time 0→1, recovery 1→2
    assert("natural recovery: counter increments to 2 after 2 periods",
        d1.recovery == 2, 2, d1.recovery)

    RLTestUtils:cleanupTestAnimals()

    -- Test recovery threshold reached (deterministic counter)
    local d2 = Disease.new(dt, false, 0)
    local animal2 = RLTestUtils:createTestAnimal("COW")
    table.insert(animal2.diseases, d2)

    d2:onPeriodChanged(animal2, false)  -- recovery 0→1
    d2:onPeriodChanged(animal2, false)  -- recovery 1→2
    d2:onPeriodChanged(animal2, false)  -- recovery 2→3 (>= type.recovery=3, eligible for cure)
    assert("natural recovery: counter reaches threshold after 3 periods",
        d2.recovery >= 3, ">=3", d2.recovery)

    RLTestUtils:cleanupTestAnimals()

    -- Statistical: type.recovery=1, run 20 fresh diseases through onPeriodChanged.
    -- Each eligible period has 75% cure chance. P(0 cures in 20) ≈ 0.0000003.
    local dtFast = RLTestUtils:createMockDiseaseType({
        recovery = 1,
        fatality = { {time = 100, value = 0} },
    })
    local curedCount = 0
    for i = 1, 20 do
        local d = Disease.new(dtFast, false, 0)
        local a = RLTestUtils:createTestAnimal("COW")
        table.insert(a.diseases, d)
        -- Two periods: first increments recovery to 1 (eligible), may cure.
        -- Second period: if not cured, recovery stays eligible.
        d:onPeriodChanged(a, false)
        d:onPeriodChanged(a, false)
        if d.cured then curedCount = curedCount + 1 end
        RLTestUtils:cleanupTestAnimals()
    end
    Log:debug("  Natural recovery statistical: %d/20 cured", curedCount)
    assert("natural recovery statistical: at least 1 cure in 20 attempts",
        curedCount >= 1, ">=1", curedCount)
end

-- =========================================================================
-- Fatality Curve
-- =========================================================================

--- Test fatality lookup using deterministic boundaries.
--- Fatality: [{time=2, value=0}, {time=100, value=1.0}]
--- time <= 2 → fatalityChance=0 (safe zone), time > 2 → fatalityChance=1.0 (guaranteed death).
--- @param assert function Assert closure from RLTestRunner
function RLDiseaseTransitionTests:testFatalityCurve(assert)
    Log:info("--- Disease Fatality Curve ---")

    local dt = RLTestUtils:createMockDiseaseType({
        fatality = { {time = 2, value = 0}, {time = 100, value = 1.0} },
        recovery = false,  -- suppress natural recovery
        treatment = false,  -- suppress treatment
    })

    -- Safe zone: time reaches 0 after 1 period (starts at -1). time=0 <= 2 → value=0, no death.
    local dSafe = Disease.new(dt, false, 0)
    local animalSafe = RLTestUtils:createTestAnimal("COW")
    table.insert(animalSafe.diseases, dSafe)
    dSafe:onPeriodChanged(animalSafe, true)  -- time -1→0
    assert("fatality safe zone: time=0 no death",
        animalSafe.isDead == false, false, animalSafe.isDead)

    RLTestUtils:cleanupTestAnimals()

    -- Death zone: advance past time=2. Need 4 periods: time goes -1→0→1→2→3.
    -- At time=3: fatality[1].time=2, 3 <= 2 = false. fatality[2] is last → value=1.0.
    -- math.random() < 1.0 is always true → guaranteed death.
    local dFatal = Disease.new(dt, false, 0)
    local animalFatal = RLTestUtils:createTestAnimal("COW")
    table.insert(animalFatal.diseases, dFatal)
    dFatal:onPeriodChanged(animalFatal, true)  -- time=0
    dFatal:onPeriodChanged(animalFatal, true)  -- time=1
    dFatal:onPeriodChanged(animalFatal, true)  -- time=2
    dFatal:onPeriodChanged(animalFatal, true)  -- time=3 → death
    assert("fatality death zone: time=3 animal dies",
        animalFatal.isDead == true, true, animalFatal.isDead)

    RLTestUtils:cleanupTestAnimals()
end
