import com.correx.core.events.events.RetryAttemptedEvent
import com.correx.core.events.execution.RetryPolicy
import com.correx.core.events.orchestration.OrchestrationState
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.kernel.retry.DefaultRetryCoordinator
import com.correx.core.kernel.retry.FailureFingerprint
import com.correx.core.kernel.retry.RetryDecision
import com.correx.infrastructure.persistence.InMemoryEventStore
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.test.currentTime
import kotlinx.coroutines.test.runTest
import org.junit.jupiter.api.Assertions
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test
import java.util.*
import kotlin.test.assertEquals

class DefaultRetryCoordinatorTest {

    private lateinit var eventStore: EventStore
    private lateinit var retryCoordinator: DefaultRetryCoordinator

    @BeforeEach
    fun setup() {
        eventStore = InMemoryEventStore()
        retryCoordinator = DefaultRetryCoordinator(eventStore)
    }

    @Test
    fun `should return true when currentAttempt is 0 and maxAttempts is 3`() = runTest {
        val sessionId = SessionId(UUID.randomUUID().toString())
        val stageId = StageId("stage1")
        val policy = RetryPolicy(maxAttempts = 3, backoffMs = 0L)

        val result = retryCoordinator.shouldRetry(
            sessionId = sessionId,
            stageId = stageId,
            currentAttempt = 0,
            policy = policy,
            failureReason = "test failure",
        )

        Assertions.assertTrue(result)

        val events = eventStore.read(sessionId)
        Assertions.assertEquals(1, events.size)
        val event = events.first().payload as RetryAttemptedEvent
        Assertions.assertEquals(1, event.attemptNumber)
        Assertions.assertEquals(stageId, event.stageId)
        Assertions.assertEquals("test failure", event.failureReason)
    }

    @Test
    fun `should return true when currentAttempt is 2 and maxAttempts is 3`() = runTest {
        val sessionId = SessionId(UUID.randomUUID().toString())
        val stageId = StageId("stage1")
        val policy = RetryPolicy(maxAttempts = 3, backoffMs = 0L)

        // First retry
        retryCoordinator.shouldRetry(
            sessionId = sessionId,
            stageId = stageId,
            currentAttempt = 0,
            policy = policy,
            failureReason = "test failure",
        )

        // Second retry (currentAttempt = 2, which is the third attempt)
        val result = retryCoordinator.shouldRetry(
            sessionId = sessionId,
            stageId = stageId,
            currentAttempt = 2,
            policy = policy,
            failureReason = "test failure",
        )

        Assertions.assertTrue(result)

        val events = eventStore.read(sessionId)
        Assertions.assertEquals(2, events.size)
        val lastEvent = events.last().payload as RetryAttemptedEvent
        Assertions.assertEquals(3, lastEvent.attemptNumber)
        Assertions.assertEquals("test failure", lastEvent.failureReason)
    }

    @Test
    fun `should return false when currentAttempt equals maxAttempts`() = runTest {
        val sessionId = SessionId(UUID.randomUUID().toString())
        val stageId = StageId("stage1")
        val policy = RetryPolicy(maxAttempts = 3, backoffMs = 0L)

        // First two retries
        retryCoordinator.shouldRetry(
            sessionId = sessionId,
            stageId = stageId,
            currentAttempt = 0,
            policy = policy,
            failureReason = "test failure",
        )
        retryCoordinator.shouldRetry(
            sessionId = sessionId,
            stageId = stageId,
            currentAttempt = 1,
            policy = policy,
            failureReason = "test failure",
        )

        // currentAttempt = 3 equals maxAttempts = 3, should not retry
        val result = retryCoordinator.shouldRetry(
            sessionId = sessionId,
            stageId = stageId,
            currentAttempt = 3,
            policy = policy,
            failureReason = "test failure",
        )

        Assertions.assertFalse(result)

        val events = eventStore.read(sessionId)
        Assertions.assertEquals(2, events.size)
    }

    @Test
    fun `should return false when currentAttempt is greater than maxAttempts`() = runTest {
        val sessionId = SessionId(UUID.randomUUID().toString())
        val stageId = StageId("stage1")
        val policy = RetryPolicy(maxAttempts = 3, backoffMs = 0L)

        // First retry
        retryCoordinator.shouldRetry(
            sessionId = sessionId,
            stageId = stageId,
            currentAttempt = 0,
            policy = policy,
            failureReason = "test failure",
        )

        // currentAttempt = 5 > maxAttempts = 3, should not retry
        val result = retryCoordinator.shouldRetry(
            sessionId = sessionId,
            stageId = stageId,
            currentAttempt = 5,
            policy = policy,
            failureReason = "test failure",
        )

        Assertions.assertFalse(result)

        val events = eventStore.read(sessionId)
        Assertions.assertEquals(1, events.size)
    }

    @Test
    @OptIn(ExperimentalCoroutinesApi::class)
    fun `should delay by backoffMs before retrying`() = runTest {
        val policy = RetryPolicy(maxAttempts = 3, backoffMs = 1000L)

        retryCoordinator.shouldRetry(
            sessionId = SessionId(UUID.randomUUID().toString()),
            stageId = StageId("stage1"),
            currentAttempt = 0,
            policy = policy,
            failureReason = "test failure",
        )

        // virtual time should have advanced by backoffMs
        assertEquals(1000L, currentTime)
    }

    // --- decide(): per-gate, progress-aware retry decision ---

    private val sessionId = SessionId(UUID.randomUUID().toString())
    private val stageId = StageId("stage1")

    @Test
    fun `decide charges the gate budget when the fingerprint is unchanged from the previous attempt`() = runTest {
        val policy = RetryPolicy(maxAttempts = 3, backoffMs = 0L)
        val state = OrchestrationState(
            gateRetryBudgets = mapOf("contract" to 1),
            gateFailureFingerprints = mapOf("contract" to FailureFingerprint.of("same-reason")),
        )

        val decision = retryCoordinator.decide(
            sessionId, stageId, gate = "contract", failureReason = "same-reason", state = state, policy = policy,
        )

        assertEquals(RetryDecision.Retry, decision)
        val event = eventStore.read(sessionId).single().payload as RetryAttemptedEvent
        assertEquals(true, event.charged)
        assertEquals(2, event.attemptNumber)
        assertEquals("contract", event.gate)
    }

    @Test
    fun `decide is free when the fingerprint changed — progress on the run`() = runTest {
        val policy = RetryPolicy(maxAttempts = 1, backoffMs = 0L)
        // Budget already at maxAttempts for this gate — a charged retry would be Exhausted, but a
        // changed fingerprint (progress) must still be free and retry.
        val state = OrchestrationState(
            gateRetryBudgets = mapOf("contract" to 1),
            gateFailureFingerprints = mapOf("contract" to "old-reason"),
        )

        val decision = retryCoordinator.decide(
            sessionId, stageId, gate = "contract", failureReason = "new-reason", state = state, policy = policy,
        )

        assertEquals(RetryDecision.Retry, decision)
        val event = eventStore.read(sessionId).single().payload as RetryAttemptedEvent
        assertEquals(false, event.charged)
        assertEquals(1, event.attemptNumber) // unchanged — this retry wasn't charged
    }

    @Test
    fun `decide returns Exhausted without emitting an event once the gate budget is spent`() = runTest {
        val policy = RetryPolicy(maxAttempts = 2, backoffMs = 0L)
        val state = OrchestrationState(
            gateRetryBudgets = mapOf("contract" to 2),
            gateFailureFingerprints = mapOf("contract" to FailureFingerprint.of("same-reason")),
        )

        val decision = retryCoordinator.decide(
            sessionId, stageId, gate = "contract", failureReason = "same-reason", state = state, policy = policy,
        )

        assertEquals(RetryDecision.Exhausted, decision)
        Assertions.assertTrue(eventStore.read(sessionId).isEmpty())
    }

    @Test
    fun `decide keys the budget by gate id — a different gate has its own fresh budget`() = runTest {
        val policy = RetryPolicy(maxAttempts = 1, backoffMs = 0L)
        val state = OrchestrationState(
            gateRetryBudgets = mapOf("contract" to 1),
            gateFailureFingerprints = mapOf("contract" to FailureFingerprint.of("same-reason")),
        )

        // "contract" is already exhausted at maxAttempts=1, but "review" has never charged.
        val decision = retryCoordinator.decide(
            sessionId, stageId, gate = "review", failureReason = "some finding", state = state, policy = policy,
        )

        assertEquals(RetryDecision.Retry, decision)
    }

    @Test
    fun `decide honours perGateMaxAttempts override`() = runTest {
        val policy = RetryPolicy(maxAttempts = 1, backoffMs = 0L, perGateMaxAttempts = mapOf("review" to 5))
        val state = OrchestrationState(
            gateRetryBudgets = mapOf("review" to 4),
            gateFailureFingerprints = mapOf("review" to "same-reason"),
        )

        val decision = retryCoordinator.decide(
            sessionId, stageId, gate = "review", failureReason = "same-reason", state = state, policy = policy,
        )

        assertEquals(RetryDecision.Retry, decision)
    }
}
