package work.slhaf.partner.core.cognition.context

import com.alibaba.fastjson2.JSONObject
import org.w3c.dom.Document
import org.w3c.dom.Element
import work.slhaf.partner.common.base.Block
import work.slhaf.partner.framework.agent.log.TraceEvent
import work.slhaf.partner.framework.agent.log.TraceRecorder
import java.time.Duration
import java.time.Instant
import java.util.*
import java.util.concurrent.locks.ReentrantReadWriteLock
import kotlin.concurrent.write
import kotlin.math.max
import kotlin.math.min

class ContextWorkspace {

    private val stateSet = mutableSetOf<ContextBlock>()
    private val lock = ReentrantReadWriteLock()

    /**
     * 根据传入的 [ContextBlock.FocusedDomain] 列表，获取上下文块
     * @param domains 需要获取上下文的域列表，顺序将决定权重优先级，按照列表排序将具备线性权重分层，最终反映到 blockContent 列表的排序上
     */
    fun resolve(domains: List<ContextBlock.FocusedDomain>): ResolvedContext = lock.write {
        if (domains.isEmpty()) {
            return@write ResolvedContext(emptyList())
        }
        val primaryDomain = domains.first()

        val domainWeights = domains
            .distinct()
            .withIndex()
            .associate { (index, domain) -> domain to (domains.size - index) }

        val activeBlocks = mutableListOf<ResolvedContextBlock>()
        val iterator = stateSet.iterator()
        while (iterator.hasNext()) {
            val block = iterator.next()
            val fadedScore = block.applyTimeFade()
            if (fadedScore <= 0.0) {
                iterator.remove()
                continue
            }

            val matchedDomains = block.focusedOn.intersect(domainWeights.keys)
            if (matchedDomains.isEmpty()) {
                continue
            }

            val exposure = if (primaryDomain in matchedDomains) {
                ContextBlock.Exposure.PRIMARY
            } else {
                ContextBlock.Exposure.SECONDARY
            }
            val activationScore = block.activate(exposure)
            if (activationScore <= 0.0) {
                iterator.remove()
                continue
            }

            activeBlocks += ResolvedContextBlock(
                block = block,
                domainWeight = matchedDomains.sumOf { domainWeights.getValue(it) },
                activationScore = activationScore,
                renderedBlock = block.render(exposure)
            )
        }

        val blocks = activeBlocks
            .sortedWith(
                compareByDescending<ResolvedContextBlock> { it.domainWeight }
                    .thenByDescending { it.activationScore }
                    .thenBy { it.block.sourceKey.blockName }
                    .thenBy { it.block.sourceKey.source }
            )
            .groupBy { it.block.sourceKey }
            .values
            .map { groupedBlocks ->
                if (groupedBlocks.size == 1) {
                    renderResolvedBlock(groupedBlocks.first())
                } else {
                    AggregatedBlockContent(groupedBlocks)
                }
            }
        ResolvedContext(blocks)
    }

    private fun renderResolvedBlock(resolved: ResolvedContextBlock): BlockContent {
        return resolved.renderedBlock
    }


    /**
     * @param contextBlock 注册的新上下文块
     */
    fun register(contextBlock: ContextBlock) = lock.write {
        val changedSourceKeys = linkedSetOf<ContextBlock.SourceKey>()
        val iterator = stateSet.iterator()
        while (iterator.hasNext()) {
            val currentBlock = iterator.next()
            if (!currentBlock.sameWith(contextBlock)) {
                continue
            }

            if (currentBlock.applyReplaceFade() <= 0.0) {
                iterator.remove()
                changedSourceKeys.add(currentBlock.sourceKey)
            }
        }
        stateSet += contextBlock
        changedSourceKeys.add(contextBlock.sourceKey)
        record("register", changedSourceKeys)
    }

    fun expire(blockName: String, source: String) = lock.write {
        val sourceKey = ContextBlock.SourceKey(blockName, source)
        val changedSourceKeys = linkedSetOf<ContextBlock.SourceKey>()
        val iterator = stateSet.iterator()
        while (iterator.hasNext()) {
            val block = iterator.next()
            if (block.sourceKey == sourceKey) {
                iterator.remove()
                changedSourceKeys.add(block.sourceKey)
            }
        }
        if (changedSourceKeys.isNotEmpty()) {
            record("expire", changedSourceKeys)
        }
    }

    private fun record(action: String, changedSourceKeys: Set<ContextBlock.SourceKey>) {
        val payload = JSONObject()
        payload["action"] = action
        payload["changedSourceKeys"] = changedSourceKeys
            .sortedWith(
                compareBy<ContextBlock.SourceKey> { it.blockName }
                    .thenBy { it.source }
            )
            .map { sourceKey ->
                JSONObject.of(
                    "blockName", sourceKey.blockName,
                    "source", sourceKey.source
                )
            }
        payload["blocks"] = stateSet
            .sortedWith(
                compareBy<ContextBlock> { it.sourceKey.blockName }
                    .thenBy { it.sourceKey.source }
            )
            .map(::blockSnapshot)
        TraceRecorder.record(TraceEvent("context-workspace", payload))
    }

    private fun blockSnapshot(block: ContextBlock): JSONObject {
        val payload = JSONObject()
        payload["blockName"] = block.sourceKey.blockName
        payload["source"] = block.sourceKey.source
        payload["focusedOn"] = block.focusedOn.map { it.name }
        payload["fullRendered"] = block.blockContent.encodeToXmlString()
        payload["compactRendered"] = block.compactBlock.encodeToXmlString()
        payload["abstractRendered"] = block.abstractBlock.encodeToXmlString()
        return payload
    }

}

private data class ResolvedContextBlock(
    val block: ContextBlock,
    val domainWeight: Int,
    val activationScore: Double,
    val renderedBlock: BlockContent
)

data class ContextBlock @JvmOverloads constructor(
    val blockContent: BlockContent,
    val compactBlock: BlockContent = blockContent,
    val abstractBlock: BlockContent = compactBlock,
    /**
     * 该 block 集中在哪些域
     */
    val focusedOn: Set<FocusedDomain>,

    /**
     * 新的 [blockContent] 属性与其相同时，发生的衰退步长
     */
    private val replaceFadeFactor: Double,

    /**
     * 随时间发生的衰退步长，按照分钟定义
     */
    private val timeFadeFactor: Double,

    /**
     * 触发一次激活时，发生的强化步长
     */
    private val activateFactor: Double
) {
    internal enum class Exposure {
        PRIMARY,
        SECONDARY
    }

    private enum class ProjectionLevel {
        ABSTRACT,
        COMPACT,
        FULL
    }

    /**
     * 默认活跃分数，降低至0时将在 [ContextWorkspace] 中移除该 block
     * 此外还参与到当同源 block 存在时的排序，按该分数升序排列，只影响同源 block 间的顺序
     */
    private var activationScore = 100.0
    private var lastTouchedAt = Instant.now()

    enum class FocusedDomain {
        ACTION,
        MEMORY,
        PERCEIVE,
        COGNITION,
        COMMUNICATION,
    }

    internal val sourceKey: SourceKey
        get() = SourceKey(blockContent.blockName, blockContent.source)

    fun applyTimeFade(): Double {
        refreshByElapsedTime()
        return activationScore
    }

    fun applyReplaceFade(): Double {
        refreshByElapsedTime()
        activationScore = max(0.0, activationScore - replaceFadeFactor)
        return activationScore
    }

    internal fun activate(exposure: Exposure): Double {
        refreshByElapsedTime()
        val currentLevel = currentProjectionLevel()
        val increasedScore = when (exposure) {
            Exposure.PRIMARY -> activationScore + when (currentLevel) {
                ProjectionLevel.FULL -> activateFactor
                ProjectionLevel.COMPACT -> activateFactor * 0.6
                ProjectionLevel.ABSTRACT -> activateFactor * 0.6
            }

            Exposure.SECONDARY -> activationScore + when (currentLevel) {
                ProjectionLevel.COMPACT -> activateFactor * 0.2
                ProjectionLevel.ABSTRACT -> activateFactor * 0.1
                ProjectionLevel.FULL -> 0.0
            }
        }
        activationScore = min(activationCeiling(exposure, currentLevel), increasedScore)
        return activationScore
    }

    private fun refreshByElapsedTime() {
        val now = Instant.now()
        val elapsedSeconds = Duration.between(lastTouchedAt, now).toMillis() / 1000.0
        activationScore = max(0.0, activationScore - elapsedSeconds * (timeFadeFactor / 60.0))
        lastTouchedAt = now
    }

    fun sameWith(contextBlock: ContextBlock): Boolean {
        return this.sourceKey == contextBlock.sourceKey
    }

    internal fun render(exposure: Exposure): BlockContent {
        return when (exposure) {
            Exposure.PRIMARY -> when (currentProjectionLevel()) {
                ProjectionLevel.FULL -> blockContent
                ProjectionLevel.COMPACT, ProjectionLevel.ABSTRACT -> compactBlock
            }

            Exposure.SECONDARY -> when (currentProjectionLevel()) {
                ProjectionLevel.ABSTRACT -> abstractBlock
                ProjectionLevel.COMPACT, ProjectionLevel.FULL -> compactBlock
            }
        }
    }

    fun render(): BlockContent {
        return when (currentProjectionLevel()) {
            ProjectionLevel.ABSTRACT -> abstractBlock
            ProjectionLevel.COMPACT -> compactBlock
            ProjectionLevel.FULL -> blockContent
        }
    }

    private fun currentProjectionLevel(): ProjectionLevel {
        return when {
            activationScore < ABSTRACT_TO_COMPACT_THRESHOLD -> ProjectionLevel.ABSTRACT
            activationScore < COMPACT_TO_FULL_THRESHOLD -> ProjectionLevel.COMPACT
            else -> ProjectionLevel.FULL
        }
    }

    private fun activationCeiling(exposure: Exposure, currentLevel: ProjectionLevel): Double {
        return when (exposure) {
            Exposure.PRIMARY -> when (currentLevel) {
                ProjectionLevel.ABSTRACT -> COMPACT_TO_FULL_THRESHOLD - PROJECTION_EPSILON
                ProjectionLevel.COMPACT, ProjectionLevel.FULL -> MAX_ACTIVATION_SCORE
            }

            Exposure.SECONDARY -> when (currentLevel) {
                ProjectionLevel.ABSTRACT -> ABSTRACT_TO_COMPACT_THRESHOLD - PROJECTION_EPSILON
                ProjectionLevel.COMPACT -> COMPACT_TO_FULL_THRESHOLD - PROJECTION_EPSILON
                ProjectionLevel.FULL -> MAX_ACTIVATION_SCORE
            }
        }
    }

    data class SourceKey(
        val blockName: String,
        val source: String
    )

    companion object {
        private const val MAX_ACTIVATION_SCORE = 100.0
        private const val ABSTRACT_TO_COMPACT_THRESHOLD = 30.0
        private const val COMPACT_TO_FULL_THRESHOLD = 70.0
        private const val PROJECTION_EPSILON = 0.000001
    }
}

private class AggregatedBlockContent(
    private val groupedBlocks: List<ResolvedContextBlock>
) : BlockContent(
    groupedBlocks.first().block.sourceKey.blockName,
    groupedBlocks.first().block.sourceKey.source,
    groupedBlocks.maxByOrNull { it.renderedBlock.urgency.ordinal }?.renderedBlock?.urgency ?: Urgency.NORMAL
) {

    override fun fillXml(document: Document, root: Element) {
        val snapshotIndex = groupedBlocks.withIndex()
            .maxWithOrNull(
                compareBy<IndexedValue<ResolvedContextBlock>> { it.value.activationScore }
                    .thenBy { it.index }
            )?.index ?: 0

        groupedBlocks.forEachIndexed { index, groupedBlock ->
            val tagName = if (index == snapshotIndex) "snapshot" else "history_snapshot"
            val wrapper = document.createElement(tagName)
            val renderedBlock = groupedBlock.renderedBlock
            val encoded = renderedBlock.encodeToXml()
            val attributes = encoded.attributes
            for (attributeIndex in 0 until attributes.length) {
                val attribute = attributes.item(attributeIndex)
                wrapper.setAttribute(attribute.nodeName, attribute.nodeValue)
            }
            root.appendChild(wrapper)

            val childNodes = encoded.childNodes
            for (childIndex in 0 until childNodes.length) {
                wrapper.appendChild(document.importNode(childNodes.item(childIndex), true))
            }
        }
    }
}

abstract class BlockContent @JvmOverloads protected constructor(
    blockName: String,
    val source: String,
    val urgency: Urgency = Urgency.NORMAL
) : Block(blockName) {

    enum class Urgency {
        LOW,
        NORMAL,
        HIGH,
        CRITICAL
    }

    override fun appendRootAttributes(): Map<String, String> {
        return mapOf(
            "source" to source,
            "urgency" to urgency.name.lowercase(Locale.ROOT)
        )
    }

}

abstract class CommunicationBlockContent @JvmOverloads constructor(
    blockName: String,
    source: String,
    urgency: Urgency = Urgency.NORMAL,
    val type: Projection = Projection.CONTEXT,
) : BlockContent(
    blockName,
    source,
    urgency
) {

    enum class Projection {
        SUPPLY,
        CONTEXT
    }
}
