package com.correx.infrastructure.inference.commons

import com.correx.core.events.types.StageId
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.NoEligibleProviderException

/**
 * An [InferenceRouter] that owns the local model lifecycle: it resolves a single target model
 * per stage and ensures it is resident (via [ModelManager.ensureLoaded]) before returning the
 * managed provider to run inference against.
 *
 * Resolution precedence (per the model-lifecycle spec):
 * ```
 * stage.modelId  >  capability match over configured models  >  defaultModelId
 * ```
 * The manual-swap pin (a higher-precedence override) layers in on top of this in a later slice.
 *
 * Because the returned provider is the live managed provider (its [InferenceProvider.id] changes
 * with the resident model), there is no stale router-level health cache to invalidate on swap —
 * selection always reflects the model the manager just made resident.
 */
class ManagedInferenceRouter(
    private val manager: ModelManager,
    descriptors: List<ModelDescriptor>,
    private val defaultModelId: String,
) : InferenceRouter {

    private val byId: Map<String, ModelDescriptor> = descriptors.associateBy { it.modelId }
    private val ordered: List<ModelDescriptor> = descriptors

    /**
     * Manual-swap pin — global operator intent. Live runtime state (decision D1): not event-sourced,
     * so a restart reverts to the default model. When set, it overrides per-stage selection.
     */
    @Volatile
    private var pinnedModelId: String? = null

    override suspend fun route(
        stageId: StageId,
        requiredCapabilities: Set<ModelCapability>,
    ): InferenceProvider = route(stageId, requiredCapabilities, null)

    override suspend fun route(
        stageId: StageId,
        requiredCapabilities: Set<ModelCapability>,
        modelId: String?,
    ): InferenceProvider {
        val targetId = pinnedModelId
            ?: modelId
            ?: capabilityMatch(requiredCapabilities)
            ?: defaultModelId
        val descriptor = byId[targetId]
            ?: throw NoEligibleProviderException(stageId, requiredCapabilities)
        return manager.ensureLoaded(descriptor)
    }

    /**
     * Manual swap from the operator: make [modelId] resident and pin it so it beats per-stage
     * selection for all later work until [clearPin]. The load emits a `ModelLoadedEvent` (when the
     * resident model actually changes), which surfaces to clients via the event stream.
     *
     * @throws NoEligibleProviderException if [modelId] is not a configured model.
     */
    suspend fun swap(modelId: String): InferenceProvider {
        val descriptor = byId[modelId]
            ?: throw NoEligibleProviderException(StageId("manual-swap"), emptySet())
        val provider = manager.ensureLoaded(descriptor)
        pinnedModelId = modelId
        return provider
    }

    /** Clear the manual-swap pin; later stages resume per-stage selection. */
    fun clearPin() {
        pinnedModelId = null
    }

    /** The currently pinned model id, or null when no manual swap is active. */
    fun pinnedModel(): String? = pinnedModelId

    /** All configured model ids, in config order — the swap candidates surfaced to clients. */
    fun availableModelIds(): List<String> = ordered.map { it.modelId }

    /** The currently resident model id, or null when none is loaded. */
    fun currentModelId(): String? = manager.currentModel()?.modelId

    /**
     * Selects the configured model whose declared capabilities cover every required capability,
     * preferring the highest summed score over those required capabilities. Returns null when no
     * capabilities are required or none of the models satisfy them (callers fall back to default).
     */
    private fun capabilityMatch(required: Set<ModelCapability>): String? {
        if (required.isEmpty()) return null
        return ordered
            .filter { descriptor ->
                val declared = descriptor.capabilities.map { it.capability }.toSet()
                declared.containsAll(required)
            }
            .maxByOrNull { descriptor ->
                descriptor.capabilities
                    .filter { it.capability in required }
                    .sumOf { it.score }
            }
            ?.modelId
    }
}
