package ru.yandex.direct.infrastructure.mysql.sharding

import org.jetbrains.exposed.sql.Database
import org.jetbrains.exposed.sql.transactions.transaction
import ru.yandex.direct.domain.client.ClientID
import ru.yandex.direct.domain.retargeting.RetargetingConditionID
import ru.yandex.direct.domain.retargeting.RetargetingID

interface ShardedDatabase {
    val ppcdict: Database
    val shards: Map<Int, Database>

    /**
     * Возвращает шард клиента (если клиенту не присвоен шард - присваивает новый).
     */
    fun shard(clientId: ClientID): Database?

    /**
     * Возвращает шард для условия ретаргетинга.
     */
    fun shard(retargetingConditionId: RetargetingConditionID): Database?

    fun generateConditionId(clientId: ClientID): RetargetingConditionID?

    fun generateRetargetingId(): RetargetingID?
}

// TODO: Build Database connection from shard number
class ShardedPPCDict(
    override val ppcdict: Database,
    override val shards: Map<Int, Database>,
) : ShardedDatabase {
    /**
     * Выбирает случайный шард для идентификатора клиента, записывает связку в [ShardClientIds].
     */
    private fun assignShardForClient(clientId: ClientID): Database? = transaction(ppcdict) {
        val client = ShardClientId.new(clientId) {
            // TODO: Implement weightedShuffle
            shard = shards.keys.random()
        }

        shards[client.shard]
    }

    /**
     * Возвращает шард клиента (если клиенту не присвоен шард - присваивает новый).
     */
    override fun shard(clientId: ClientID): Database? = transaction(ppcdict) {
        when (val client = ShardClientId.findById(clientId)) {
            is ShardClientId -> shards[client.shard]
            else -> assignShardForClient(clientId)
        }
    }

    /**
     * Возвращает шард для условия ретаргетинга.
     */
    override fun shard(retargetingConditionId: RetargetingConditionID): Database? = transaction(ppcdict) {
        shards[ShardIncRetCondId.findById(retargetingConditionId)?.client?.shard]
    }

    override fun generateConditionId(clientId: ClientID): RetargetingConditionID? = transaction(ppcdict) {
        // TODO: For optimizations ShardClientIDs.batchInsert can be used
        when (val c = ShardClientId.findById(clientId)) {
            is ShardClientId -> ShardIncRetCondId.new { client = c }.id.value
            else -> null
        }
    }

    override fun generateRetargetingId(): RetargetingID = transaction(ppcdict) {
        IncRetId.new{}.id.value
    }
}
