package com.correx.core.inference

import com.correx.core.events.types.ProviderId
import com.correx.core.events.types.StageId
import kotlinx.coroutines.delay
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlin.time.Duration
import kotlin.time.Duration.Companion.seconds
import kotlin.time.TimeMark
import kotlin.time.TimeSource

private val log = org.slf4j.LoggerFactory.getLogger(DefaultInferenceRouter::class.java)

private data class HealthEntry(
    val health: ProviderHealth,
    val mark: TimeMark,
)

class DefaultInferenceRouter(
    private val registry: ProviderRegistry,
    private val strategy: RoutingStrategy,
    private val cacheTtl: Duration = 5.seconds,
    private val timeSource: TimeSource = TimeSource.Monotonic,
    // A provider that briefly drops (crash + qa-stack restart) shouldn't collapse into a hard
    // NoEligibleProvider abort — give it bounded time to come back before giving up. This only
    // applies when the capability IS configured on some provider but that provider is currently
    // unhealthy; a capability nobody ever declared fails immediately (see routeCapabilityCandidates).
    private val unavailableRetryAttempts: Int = 3,
    private val unavailableRetryDelay: Duration = 2.seconds,
) : InferenceRouter {

    private val cache = mutableMapOf<ProviderId, HealthEntry>()
    private val locks = mutableMapOf<ProviderId, Mutex>()
    private val mapMutex = Mutex()

    private suspend fun lockFor(id: ProviderId): Mutex =
        mapMutex.withLock { locks.getOrPut(id) { Mutex() } }

    private suspend fun cachedHealth(provider: InferenceProvider): ProviderHealth {
        val existing = mapMutex.withLock { cache[provider.id] }
        if (existing != null && existing.mark.elapsedNow() < cacheTtl) return existing.health

        return lockFor(provider.id).withLock {
            // re-check after acquiring lock — another coroutine may have refreshed
            val recheck = mapMutex.withLock { cache[provider.id] }
            if (recheck != null && recheck.mark.elapsedNow() < cacheTtl) return@withLock recheck.health

            val fresh = provider.healthCheck()
            mapMutex.withLock { cache[provider.id] = HealthEntry(fresh, timeSource.markNow()) }
            fresh
        }
    }

    private suspend fun refreshedHealth(provider: InferenceProvider): ProviderHealth =
        lockFor(provider.id).withLock {
            val fresh = provider.healthCheck()
            mapMutex.withLock { cache[provider.id] = HealthEntry(fresh, timeSource.markNow()) }
            fresh
        }

    // Bypasses healthCheck()/TTL entirely — writes Unavailable straight into the cache so the very
    // next route() call gates on it, closing the ~18s reactive-poll lag (#300). A later cache-TTL
    // expiry or the bounded-wait re-check in route() will naturally pick the provider back up once
    // its own healthCheck() reports healthy again.
    override suspend fun reportFailure(providerId: ProviderId, reason: String) {
        lockFor(providerId).withLock {
            mapMutex.withLock { cache[providerId] = HealthEntry(ProviderHealth.Unavailable(reason), timeSource.markNow()) }
        }
    }

    override suspend fun route(stageId: StageId, requiredCapabilities: Set<ModelCapability>): InferenceProvider {
        val candidates = requiredCapabilities
            .flatMap { registry.resolve(it) }
            .distinctBy { it.id }
            .ifEmpty { registry.listAll() }

        // Nobody was ever configured with this capability set — no amount of waiting fixes that,
        // fail fast instead of burning the bounded-wait budget below.
        if (requiredCapabilities.isNotEmpty() &&
            candidates.none { it.capabilities().map { c -> c.capability }.toSet().containsAll(requiredCapabilities) }
        ) {
            throw NoEligibleProviderException(stageId, requiredCapabilities)
        }

        var healthy = candidates.filter { cachedHealth(it) !is ProviderHealth.Unavailable }
        var attempt = 0
        while (healthy.isEmpty() && attempt < unavailableRetryAttempts) {
            attempt++
            log.warn(
                "route: capability {} configured but all candidates unhealthy for stage={}" +
                    " — waiting {} (attempt {}/{}) before declaring NoEligibleProvider",
                requiredCapabilities, stageId.value, unavailableRetryDelay, attempt, unavailableRetryAttempts,
            )
            delay(unavailableRetryDelay)
            healthy = candidates.filter { refreshedHealth(it) !is ProviderHealth.Unavailable }
        }
        val selected = strategy.select(healthy, requiredCapabilities)
        // Post-selection re-check closes the TOCTOU window between initial filter and dispatch.
        when (val postHealth = selected.healthCheck()) {
            is ProviderHealth.Unavailable -> throw ProviderUnavailableException(selected.id, postHealth.reason)
            else -> Unit
        }
        return selected
    }

    override suspend fun route(
        stageId: StageId,
        requiredCapabilities: Set<ModelCapability>,
        modelId: String?,
    ): InferenceProvider {
        if (modelId == null) return route(stageId, requiredCapabilities)

        val all = registry.listAll()
        val healthy = all.filter { cachedHealth(it) !is ProviderHealth.Unavailable }
        // Registry ids are adapter-prefixed ("llama-cpp:narrator") while the operator's TOML
        // declares the bare id ("narrator") — accept either form.
        val matched = healthy.firstOrNull {
            it.id.value == modelId || it.id.value.substringAfter(':') == modelId
        }
        if (matched == null) {
            log.warn(
                "narration model_id '{}' matched no healthy provider — falling back to capability routing",
                modelId,
            )
            return route(stageId, requiredCapabilities)
        }
        when (val postHealth = matched.healthCheck()) {
            is ProviderHealth.Unavailable -> {
                log.warn(
                    "narration model_id '{}' provider failed post-select health check ({})" +
                        " — falling back to capability routing",
                    modelId,
                    postHealth.reason,
                )
                return route(stageId, requiredCapabilities)
            }
            else -> Unit
        }
        return matched
    }
}
