package work.slhaf.partner.module.action.scheduler

import com.cronutils.model.CronType
import com.cronutils.model.definition.CronDefinition
import com.cronutils.model.definition.CronDefinitionBuilder
import com.cronutils.model.time.ExecutionTime
import com.cronutils.parser.CronParser
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import org.slf4j.LoggerFactory
import work.slhaf.partner.core.action.ActionCapability
import work.slhaf.partner.core.action.entity.Action
import work.slhaf.partner.core.action.entity.Schedulable
import work.slhaf.partner.core.action.entity.SchedulableExecutableAction
import work.slhaf.partner.framework.agent.config.ConfigCenter
import work.slhaf.partner.framework.agent.factory.capability.annotation.InjectCapability
import work.slhaf.partner.framework.agent.factory.component.abstracts.AbstractAgentModule
import work.slhaf.partner.framework.agent.factory.component.annotation.Init
import work.slhaf.partner.framework.agent.factory.component.annotation.InjectModule
import work.slhaf.partner.module.action.executor.ActionExecutor
import java.io.Closeable
import java.nio.file.Path
import java.time.Duration
import java.time.ZonedDateTime
import java.time.temporal.ChronoUnit
import java.util.*
import java.util.stream.Collectors
import kotlin.jvm.optionals.getOrNull
import kotlin.time.Duration.Companion.milliseconds

class ActionScheduler : AbstractAgentModule.Standalone() {

    @InjectCapability
    private lateinit var actionCapability: ActionCapability

    @InjectModule
    private lateinit var actionExecutor: ActionExecutor

    private lateinit var timeWheel: TimeWheel
    private val runtimeSchedulables: MutableSet<Schedulable> =
        Collections.synchronizedSet(mutableSetOf())
    private val schedulerScope =
        CoroutineScope(Dispatchers.Default + SupervisorJob() + CoroutineName("ActionScheduler"))

    companion object {
        private val log = LoggerFactory.getLogger(ActionScheduler::class.java)
    }

    @Init
    fun init() {
        fun loadScheduledActions() {
            val listScheduledActions: () -> Set<Schedulable> = {
                val persistedExecutable = actionCapability.listActions(null, null)
                    .stream()
                    .filter { it is SchedulableExecutableAction }
                    .map { it as SchedulableExecutableAction }
                    .collect(Collectors.toSet())
                val persisted: MutableSet<Schedulable> = mutableSetOf()
                persisted.addAll(persistedExecutable)
                synchronized(runtimeSchedulables) {
                    persisted.addAll(runtimeSchedulables.filter { it.enabled })
                }
                persisted
            }
            val onTrigger: (Set<Schedulable>) -> Unit = { schedulableSet ->
                schedulableSet.filterIsInstance<Action>()
                    .forEach { actionExecutor.execute(it) }
            }
            val doneCondition: (Schedulable) -> Boolean = { schedulable ->
                if (schedulable is Action) {
                    schedulable.status == Action.Status.FAILED || schedulable.status == Action.Status.SUCCESS
                } else {
                    true
                }
            }
            timeWheel = TimeWheel(listScheduledActions, onTrigger, doneCondition)
        }
        loadScheduledActions()
        setupShutdownHook()
    }

    private fun setupShutdownHook() {
        Runtime.getRuntime().addShutdownHook(Thread {
            timeWheel.close()
            schedulerScope.cancel()
        })
    }

    fun <T> schedule(schedulableAction: T) where T : Action, T : Schedulable = schedulerScope.launch {
        if (!schedulableAction.enabled) {
            return@launch
        }
        runtimeSchedulables.add(schedulableAction)
        log.debug("New data to schedule: {}", schedulableAction)
        timeWheel.schedule(schedulableAction)
        if (schedulableAction is SchedulableExecutableAction) {
            actionCapability.putAction(schedulableAction)
        }
    }

    fun cancel(actionId: String): Boolean {
        val runtimeMatches = synchronized(runtimeSchedulables) {
            runtimeSchedulables.filter { it.uuid == actionId }.toSet()
        }
        val persistedMatches = actionCapability.listActions(null, null)
            .asSequence()
            .filterIsInstance<SchedulableExecutableAction>()
            .filter { it.uuid == actionId }
            .toSet()

        val matches = LinkedHashSet<Schedulable>()
        matches.addAll(runtimeMatches)
        matches.addAll(persistedMatches)
        matches.forEach { it.enabled = false }

        val removedFromWheel = if (::timeWheel.isInitialized) {
            runBlocking { timeWheel.cancel(actionId) }
        } else {
            false
        }
        return matches.isNotEmpty() || removedFromWheel
    }

    private class TimeWheel(
        val listSource: () -> Set<Schedulable>,
        val onTrigger: (toTrigger: Set<Schedulable>) -> Unit,
        val doneCondition: (schedulable: Schedulable) -> Boolean
    ) : Closeable {
        private val tracePath: Path = ConfigCenter.paths.stateDir
            .resolve("trace")
            .resolve("time-wheel")
            .normalize()
            .toAbsolutePath()

        private val schedulableGroupByHour = Array<MutableSet<Schedulable>>(24) { mutableSetOf() }
        private val wheel = Array<MutableSet<Schedulable>>(60 * 60) { mutableSetOf() }
        private var recordHour: Int = -1
        private var recordDay: Int = -1
        private val state = MutableStateFlow(WheelState.SLEEPING)
        private val wheelActionsLock = Mutex()
        private val timeWheelScope = CoroutineScope(SupervisorJob() + Dispatchers.Default + CoroutineName("TimeWheel"))
        private val cronDefinition: CronDefinition = CronDefinitionBuilder.instanceDefinitionFor(CronType.QUARTZ)
        private val cronParser: CronParser = CronParser(cronDefinition)

        init {
            // 启动时间轮
            wheel()
        }

        suspend fun schedule(schedulableData: Schedulable) {
            checkThenExecute {
                val parseToZonedDateTime = parseToZonedDateTime(
                    schedulableData.scheduleType,
                    schedulableData.scheduleContent,
                    it
                ) ?: run {
                    logFailedStatus(schedulableData)
                    return@checkThenExecute
                }
                log.debug("${schedulableData.uuid} next execution time: {}", parseToZonedDateTime)
                val hour = parseToZonedDateTime.hour
                schedulableGroupByHour[hour].add(schedulableData)
                if (it.hour == hour) {
                    val wheelOffset = parseToZonedDateTime.minute * 60 + parseToZonedDateTime.second
                    wheel[wheelOffset].add(schedulableData)
                    state.value = WheelState.ACTIVE
                    log.debug("${schedulableData.uuid} scheduled at wheel offset {}", wheelOffset)
                }
            }
        }

        suspend fun cancel(actionId: String): Boolean = wheelActionsLock.withLock {
            var found = false
            for (bucket in schedulableGroupByHour) {
                var bucketFound = false
                bucket.removeIf {
                    if (it.uuid == actionId) {
                        it.enabled = false
                        bucketFound = true
                        true
                    } else {
                        false
                    }
                }
                found = found || bucketFound
            }
            for (bucket in wheel) {
                var bucketFound = false
                bucket.removeIf {
                    if (it.uuid == actionId) {
                        it.enabled = false
                        bucketFound = true
                        true
                    } else {
                        false
                    }
                }
                found = found || bucketFound
            }
            found
        }

        private fun wheel() {
            data class WheelStepResult(
                val toTrigger: Set<Schedulable>?,
                val shouldBreak: Boolean
            )

            fun collectToTrigger(
                tick: Int,
                previousTick: Int,
                triggerHour: Int,
                includeRemainingHour: Boolean = false
            ): Set<Schedulable>? {
                if (tick >= previousTick) {
                    val toTrigger = mutableSetOf<Schedulable>()
                    for (i in previousTick..tick) {
                        val bucket = wheel[i]
                        if (bucket.isNotEmpty()) {
                            toTrigger.addAll(bucket.filter { it.enabled })
                            val bucketUuids = bucket.asSequence().map { it.uuid }.toHashSet()
                            schedulableGroupByHour[triggerHour].removeIf { it.uuid in bucketUuids }
                            bucket.clear() // 避免重复触发
                        }
                    }
                    if (includeRemainingHour) {
                        val remainingHourBucket = schedulableGroupByHour[triggerHour]
                        if (remainingHourBucket.isNotEmpty()) {
                            toTrigger.addAll(remainingHourBucket.filter { it.enabled })
                            remainingHourBucket.clear()
                        }
                    }
                    return toTrigger
                }
                return null
            }

            fun handleToTrigger(toTrigger: Set<Schedulable>) {
                timeWheelScope.launch {
                    onTrigger(toTrigger)
                }
                for (schedulable in toTrigger) timeWheelScope.launch {
                    if (schedulable.scheduleType == Schedulable.ScheduleType.ONCE) {
                        return@launch
                    }

                    withTimeoutOrNull(schedulable.timeout) {
                        while (!doneCondition(schedulable)) {
                            delay(50.milliseconds)
                        }
                    }

                    if (!schedulable.enabled) {
                        return@launch
                    }

                    this@TimeWheel.schedule(schedulable)
                }
            }

            suspend fun CoroutineScope.wheel(launchingTime: ZonedDateTime, primaryTickAdvanceTime: Long) {
                val launchingHour = launchingTime.hour
                var tick = launchingTime.minute * 60 + launchingTime.second
                // 让节拍器从“启动时刻的下一秒”开始（避免立即 step=0）
                var nextTickNanos = primaryTickAdvanceTime + 1_000_000_000L
                while (isActive) {
                    // 1) 计算落后多少秒：至少 1（正常推进），也可能 >1（追赶）
                    val now0 = System.nanoTime()
                    val lagNanos = now0 - nextTickNanos
                    val step = if (lagNanos < 0) 1 else (lagNanos / 1_000_000_000L).toInt() + 1
                    val previousTick = tick
                    tick = (tick + step).coerceAtMost(wheel.lastIndex)
                    // 2) 推进节拍器：按“理论秒”前进 step 次
                    nextTickNanos += step.toLong() * 1_000_000_000L
                    val stepResult = run {
                        var shouldBreak = false
                        var toTrigger: Set<Schedulable>? = null
                        checkThenExecute(false) {
                            if (it.hour != launchingHour) {
                                shouldBreak = true
                                toTrigger = collectToTrigger(
                                    wheel.lastIndex,
                                    previousTick,
                                    launchingHour,
                                    includeRemainingHour = true
                                )
                                log.debug(
                                    "Hour changed, previousTick: {}, tick: {}, toTriggerSize: {}",
                                    previousTick,
                                    tick,
                                    toTrigger?.size
                                )
                                return@checkThenExecute
                            }
                            toTrigger = collectToTrigger(
                                tick,
                                previousTick,
                                launchingHour,
                                includeRemainingHour = tick >= wheel.lastIndex
                            )
                            if (tick >= wheel.lastIndex) {
                                state.value = WheelState.SLEEPING
                                shouldBreak = true
                            }
                        }
                        WheelStepResult(toTrigger, shouldBreak)
                    }
                    stepResult.toTrigger?.let { handleToTrigger(it) }
                    if (stepResult.shouldBreak) {
                        log.debug("Wheel stopped at tick {}", tick)
                        break
                    }
                    // 3) 精确睡到下一次理论 tick（用最新 nanoTime）
                    val now1 = System.nanoTime()
                    val sleepNanos = nextTickNanos - now1
                    if (sleepNanos > 0) {
                        delay((sleepNanos / 1_000_000L).milliseconds) // 毫秒级 delay 足够；剩余 nanos 不必忙等
                    }
                }
            }

            suspend fun wait(currentTime: ZonedDateTime) {
                val nextHour = currentTime.truncatedTo(ChronoUnit.HOURS).plusHours(1)
                val seconds = Duration.between(
                    currentTime, nextHour
                ).toMillis()
                // withTimeoutOrNull 内部已处理 seconds 小于 0 的情况
                log.debug("Start waiting {} ms at {}, target time: {}", seconds, currentTime, nextHour)
                withTimeoutOrNull(seconds.milliseconds) {
                    state.first { it == WheelState.ACTIVE }
                }
                log.debug("Waiting ended at {}", ZonedDateTime.now())
            }
            timeWheelScope.launch {
                while (isActive) {
                    // 判断是否该步入下一小时
                    var shouldWait: Boolean? = null
                    var currentTime: ZonedDateTime? = null
                    var primaryTickAdvanceTime: Long? = null
                    checkThenExecute {
                        currentTime = it
                        shouldWait = schedulableGroupByHour[it.hour].isEmpty()
                        // 由于 wheel 的启动时间可能存在延迟，而时内推进由 nanoTime 保证不会漏发，
                        // 正常的时序结束又由 tick 是否触顶、当前时是否存在额外任务触发，
                        // 而启动时无触发保障，此时一并初始化 tick 推进时间，足以应对 check 与 wheel 间的这段时间间隔
                        primaryTickAdvanceTime = System.nanoTime()
                    }
                    // 如果该时无任务则等待，插入事件可提前唤醒
                    if (shouldWait!!) {
                        // 计算距离下一小时的时间，等待
                        currentTime?.let { wait(it) }
                        continue
                    }
                    // 唤醒进行时间轮循环
                    wheel(currentTime!!, primaryTickAdvanceTime!!)
                }
            }
        }

        suspend fun checkThenExecute(finallyToExecute: Boolean = true, then: (currentTime: ZonedDateTime) -> Unit) =
            wheelActionsLock.withLock {
                fun loadActions(
                    source: Set<Schedulable>,
                    now: ZonedDateTime,
                    load: (latestExecutingTime: ZonedDateTime, schedulableData: Schedulable) -> Unit,
                    repair: () -> Unit
                ) {
                    val runLoading = {
                        for (schedulableData in source) {
                            val nextExecutingTime =
                                parseToZonedDateTime(
                                    schedulableData.scheduleType,
                                    schedulableData.scheduleContent,
                                    now
                                ) ?: run {
                                    logFailedStatus(schedulableData)
                                    continue
                                }
                            load(nextExecutingTime, schedulableData)
                        }
                    }
                    repair()
                    runLoading()
                }

                fun loadHourActions(currentTime: ZonedDateTime) {
                    val load: (ZonedDateTime, Schedulable) -> Unit =
                        { latestExecutionTime, schedulableData ->
                            val secondsTime = latestExecutionTime.minute * 60 + latestExecutionTime.second
                            wheel[secondsTime].add(schedulableData)
                            log.debug("Action loaded to hour: {}", schedulableData)
                        }
                    val repair: () -> Unit = {
                        for (set in wheel) {
                            set.clear()
                        }
                    }
                    loadActions(schedulableGroupByHour[currentTime.hour], currentTime, load, repair)
                }

                fun loadDayActions(currentTime: ZonedDateTime) {
                    val load: (ZonedDateTime, Schedulable) -> Unit =
                        { latestExecutingTime, schedulableData ->
                            schedulableGroupByHour[latestExecutingTime.hour].add(schedulableData)
                            log.debug("Action loaded to day: {}", schedulableData)
                        }
                    val repair: () -> Unit = {
                        for (set in schedulableGroupByHour) {
                            set.clear()
                        }
                    }
                    loadActions(listSource(), currentTime, load, repair)
                }

                fun refreshIfNeeded(now: ZonedDateTime) {
                    val d = now.dayOfMonth
                    val h = now.hour
                    if (d != recordDay) {
                        recordDay = d
                        recordHour = h
                        loadDayActions(now)
                        loadHourActions(now)
                    } else if (h != recordHour) {
                        recordHour = h
                        loadHourActions(now)
                    }
                }

                val now = ZonedDateTime.now()
                if (finallyToExecute) {
                    refreshIfNeeded(now)
                    then(now)
                } else {
                    then(now)
                    refreshIfNeeded(now)
                }
            }

        private fun parseToZonedDateTime(
            scheduleType: Schedulable.ScheduleType,
            scheduleContent: String,
            now: ZonedDateTime
        ): ZonedDateTime? {
            return when (scheduleType) {
                Schedulable.ScheduleType.CYCLE
                    -> {
                    val cron = try {
                        cronParser.parse(scheduleContent).validate()
                    } catch (_: Exception) {
                        return null
                    }
                    val executionTime = ExecutionTime.forCron(cron)
                    executionTime.nextExecution(now).getOrNull()
                }

                Schedulable.ScheduleType.ONCE -> {
                    val executionTime = try {
                        ZonedDateTime.parse(scheduleContent)
                    } catch (_: Exception) {
                        return null
                    }
                    if (executionTime.plusSeconds(1).isBefore(now) || executionTime.dayOfMonth != now.dayOfMonth)
                        null
                    else
                        executionTime
                }
            }
        }

        private fun logFailedStatus(scheduleData: Schedulable) {
            log.warn(
                "行动未加载，scheduleType: {}, scheduleContent: {}",
                scheduleData.scheduleType,
                scheduleData.scheduleContent,
            )
        }

        override fun close() {
            timeWheelScope.cancel()
        }

        private enum class WheelState {
            ACTIVE,
            SLEEPING,
        }
    }
}
