package com.correx.infrastructure.inference

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

/**
 * Like [FirstAvailableRoutingStrategy] it hard-filters to providers that declare every
 * [requiredCapabilities] — but among those it picks the one with the highest summed capability
 * **score** for the required set (falling back to total declared score when nothing is required).
 * Ties keep list order, so behaviour is identical to first-available when scores are equal.
 *
 * This is the seam the orchestrator relies on to steer tool-heavy stages (which request
 * `ModelCapability.ToolCalling`) to the model that calls tools best, rather than whichever healthy
 * provider happens to come first. The scores come from each provider's declared
 * `capabilities()` (config-supplied); observed per-model reliability (the `/metrics/tool-reliability`
 * groundwork) can later feed in as an additional weight.
 */
class CapabilityAwareRoutingStrategy : RoutingStrategy {
    override fun select(
        candidates: List<InferenceProvider>,
        requiredCapabilities: Set<ModelCapability>,
    ): InferenceProvider {
        val eligible = candidates.filter { provider ->
            provider.capabilities().map { it.capability }.toSet().containsAll(requiredCapabilities)
        }
        // maxByOrNull keeps the first element on ties, preserving first-available order.
        return eligible.maxByOrNull { score(it, requiredCapabilities) }
            ?: throw NoEligibleProviderException(StageId("routing"), requiredCapabilities)
    }

    private fun score(provider: InferenceProvider, required: Set<ModelCapability>): Double {
        val scores = provider.capabilities().associate { it.capability to it.score }
        return if (required.isEmpty()) scores.values.sum() else required.sumOf { scores[it] ?: 0.0 }
    }
}
