package space.nstart.pcp.pcp_satellites_service.service

import org.slf4j.LoggerFactory
import org.springframework.stereotype.Service
import java.sql.Connection
import java.util.concurrent.ConcurrentHashMap
import javax.sql.DataSource
import kotlin.system.measureTimeMillis

@Service
class ComplexPlanExecutionLockService(
    private val dataSource: DataSource
) {

    companion object {
        private const val LOCK_NAMESPACE = 4_204_202_411L
        private val fallbackLocks = ConcurrentHashMap.newKeySet<Long>()
    }

    private val logger = LoggerFactory.getLogger(this::class.java)

    fun tryAcquire(satelliteIds: Collection<Long>): ExecutionLockHandle? {
        val normalizedSatelliteIds = satelliteIds.distinct().sorted()
        if (normalizedSatelliteIds.isEmpty()) {
            logger.info("Complex plan execution lock skipped: empty satellite list")
            return ExecutionLockHandle {}
        }

        var handle: ExecutionLockHandle? = null
        val durationMs = measureTimeMillis {
            val connection = dataSource.connection
            handle = try {
                val postgres = connection.metaData.databaseProductName.contains("PostgreSQL", ignoreCase = true)
                if (postgres) {
                    acquirePostgresLocks(connection, normalizedSatelliteIds)
                } else {
                    connection.close()
                    acquireFallbackLocks(normalizedSatelliteIds)
                }
            } catch (ex: Exception) {
                runCatching { connection.close() }
                throw ex
            }
        }
        logger.info(
            "Complex plan execution lock acquire completed: satellites={}, acquired={}, durationMs={}",
            normalizedSatelliteIds,
            handle != null,
            durationMs
        )
        return handle
    }

    private fun acquirePostgresLocks(connection: Connection, satelliteIds: List<Long>): ExecutionLockHandle? {
        val acquiredKeys = mutableListOf<Long>()
        try {
            satelliteIds.forEach { satelliteId ->
                val lockKey = advisoryKey(satelliteId)
                if (!tryAcquirePostgresLock(connection, lockKey)) {
                    releasePostgresLocks(connection, acquiredKeys.asReversed())
                    connection.close()
                    logger.info(
                        "Complex plan execution lock busy: satelliteId={}, acquiredSatelliteIds={}",
                        satelliteId,
                        acquiredKeys.map(::satelliteIdFromKey)
                    )
                    return null
                }
                acquiredKeys += lockKey
            }
        } catch (ex: Exception) {
            releasePostgresLocks(connection, acquiredKeys.asReversed())
            connection.close()
            throw ex
        }

        return ExecutionLockHandle {
            releasePostgresLocks(connection, acquiredKeys.asReversed())
            connection.close()
        }
    }

    private fun acquireFallbackLocks(satelliteIds: List<Long>): ExecutionLockHandle? {
        val acquiredKeys = mutableListOf<Long>()
        satelliteIds.forEach { satelliteId ->
            val lockKey = advisoryKey(satelliteId)
            if (!fallbackLocks.add(lockKey)) {
                acquiredKeys.forEach(fallbackLocks::remove)
                logger.info(
                    "Complex plan execution fallback lock busy: satelliteId={}, acquiredSatelliteIds={}",
                    satelliteId,
                    acquiredKeys.map(::satelliteIdFromKey)
                )
                return null
            }
            acquiredKeys += lockKey
        }

        return ExecutionLockHandle {
            acquiredKeys.forEach(fallbackLocks::remove)
        }
    }

    private fun tryAcquirePostgresLock(connection: Connection, lockKey: Long): Boolean =
        connection.prepareStatement("select pg_try_advisory_lock(?)").use { statement ->
            statement.setLong(1, lockKey)
            statement.executeQuery().use { resultSet ->
                resultSet.next()
                resultSet.getBoolean(1)
            }
        }

    private fun releasePostgresLocks(connection: Connection, lockKeys: List<Long>) {
        lockKeys.forEach { lockKey ->
            runCatching {
                connection.prepareStatement("select pg_advisory_unlock(?)").use { statement ->
                    statement.setLong(1, lockKey)
                    statement.executeQuery().use { resultSet -> resultSet.next() }
                }
            }.onFailure { ex ->
                logger.warn("Failed to release complex plan advisory lock: key={}", lockKey, ex)
            }
        }
    }

    private fun advisoryKey(satelliteId: Long): Long = LOCK_NAMESPACE * 1_000_003L + satelliteId

    private fun satelliteIdFromKey(lockKey: Long): Long = lockKey - LOCK_NAMESPACE * 1_000_003L

    fun interface ExecutionLockHandle : AutoCloseable {
        override fun close()
    }
}
