package cc.unitmesh.viewer.web.webedit

import kotlinx.serialization.Serializable

/**
 * Accessibility Tree (AOM) representation for LLM Agents
 *
 * Based on: https://arxiv.org/html/2508.04412v2
 *
 * The Accessibility Tree is a semantic tree generated by browsers for screen readers.
 * It naturally ignores visual decorations and preserves only the interactive logic structure.
 *
 * Key properties:
 * - Role: The semantic role of the element (button, link, textbox, etc.)
 * - Name: The accessible name (computed from aria-label, text content, etc.)
 * - State: Interactive states (disabled, focused, checked, expanded, etc.)
 */
object AccessibilityTree {

    /**
     * Convert a DOMElement tree to AccessibilityTree format
     */
    fun fromDOM(element: DOMElement): AccessibilityNode? {
        return convertToAccessibilityNode(element, 0)
    }

    private fun convertToAccessibilityNode(element: DOMElement, depth: Int): AccessibilityNode? {
        if (depth > 15) return null

        val tagName = element.tagName.lowercase()
        val role = computeRole(tagName, element.attributes)
        val name = computeAccessibleName(element)
        val state = computeState(element.attributes)

        // Skip non-semantic container nodes without role
        if (role == null && name.isNullOrBlank() && state.isEmpty()) {
            // But process children - they might have semantic meaning
            val children = element.children.mapNotNull { convertToAccessibilityNode(it, depth + 1) }
            return if (children.size == 1) {
                children.first()
            } else if (children.isNotEmpty()) {
                // Create a generic container only if multiple meaningful children
                AccessibilityNode(
                    role = "group",
                    name = null,
                    state = emptyMap(),
                    value = null,
                    selector = element.selector,
                    children = children
                )
            } else {
                null
            }
        }

        val children = element.children.mapNotNull { convertToAccessibilityNode(it, depth + 1) }

        return AccessibilityNode(
            role = role ?: inferRoleFromTag(tagName),
            name = name,
            state = state,
            value = computeValue(element),
            selector = element.selector,
            children = children,
            description = element.attributes["aria-describedby"]?.let { 
                element.attributes["title"] 
            } ?: element.attributes["title"]
        )
    }

    /**
     * Compute ARIA role from explicit role attribute or implicit tag semantics
     */
    private fun computeRole(tagName: String, attributes: Map<String, String>): String? {
        // Explicit role takes precedence
        attributes["role"]?.let { return it }

        // Implicit roles based on HTML5 semantics
        return when (tagName) {
            "button" -> "button"
            "a" -> if (attributes.containsKey("href")) "link" else null
            "input" -> computeInputRole(attributes["type"])
            "select" -> if (attributes["multiple"] != null) "listbox" else "combobox"
            "textarea" -> "textbox"
            "img" -> "img"
            "nav" -> "navigation"
            "main" -> "main"
            "header" -> "banner"
            "footer" -> "contentinfo"
            "aside" -> "complementary"
            "article" -> "article"
            "section" -> if (attributes.containsKey("aria-label") || attributes.containsKey("aria-labelledby")) "region" else null
            "form" -> "form"
            "table" -> "table"
            "tr" -> "row"
            "th" -> "columnheader"
            "td" -> "cell"
            "ul", "ol" -> "list"
            "li" -> "listitem"
            "dialog" -> "dialog"
            "menu" -> "menu"
            "menuitem" -> "menuitem"
            "details" -> "group"
            "summary" -> "button"
            "h1", "h2", "h3", "h4", "h5", "h6" -> "heading"
            "progress" -> "progressbar"
            "meter" -> "meter"
            "option" -> "option"
            "optgroup" -> "group"
            "fieldset" -> "group"
            "legend" -> null // Part of fieldset, name computed differently
            "label" -> null // Labels provide names, not roles
            "figure" -> "figure"
            "figcaption" -> null // Part of figure
            "blockquote" -> "blockquote"
            "code", "pre" -> "code"
            "search" -> "search"
            "output" -> "status"
            else -> null
        }
    }

    /**
     * Compute input role based on type attribute
     */
    private fun computeInputRole(type: String?): String {
        return when (type?.lowercase()) {
            "button", "submit", "reset", "image" -> "button"
            "checkbox" -> "checkbox"
            "radio" -> "radio"
            "range" -> "slider"
            "number", "spinbutton" -> "spinbutton"
            "search" -> "searchbox"
            "email", "tel", "url", "password", "text", null -> "textbox"
            "file" -> "button" // File inputs act like buttons
            "hidden" -> "none"
            "color" -> "button"
            "date", "datetime-local", "month", "week", "time" -> "textbox"
            else -> "textbox"
        }
    }

    /**
     * Infer a generic role from tag name for unknown elements
     */
    private fun inferRoleFromTag(tagName: String): String {
        return when {
            tagName.startsWith("h") && tagName.length == 2 -> "heading"
            tagName in setOf("div", "span", "p") -> "generic"
            else -> "generic"
        }
    }

    /**
     * Compute accessible name following ARIA naming computation algorithm
     * Priority: aria-labelledby > aria-label > associated label > text content > title > placeholder
     */
    private fun computeAccessibleName(element: DOMElement): String? {
        val attrs = element.attributes

        // 1. aria-label takes precedence
        attrs["aria-label"]?.takeIf { it.isNotBlank() }?.let { return it }

        // 2. aria-labelledby would need DOM lookup (simplified here)
        // In a real implementation, we'd resolve the ID references

        // 3. For inputs, check placeholder
        if (element.tagName.lowercase() == "input") {
            attrs["placeholder"]?.takeIf { it.isNotBlank() }?.let { return it }
        }

        // 4. For images, use alt text
        if (element.tagName.lowercase() == "img") {
            attrs["alt"]?.let { return it }
        }

        // 5. For links/buttons, use text content
        if (element.tagName.lowercase() in setOf("a", "button", "label", "h1", "h2", "h3", "h4", "h5", "h6", "option", "summary")) {
            element.textContent?.trim()?.takeIf { it.isNotBlank() && it.length < 200 }?.let { return it }
        }

        // 6. Title attribute as fallback
        attrs["title"]?.takeIf { it.isNotBlank() }?.let { return it }

        // 7. Value for inputs
        if (element.tagName.lowercase() in setOf("input", "textarea")) {
            attrs["value"]?.takeIf { it.isNotBlank() }?.let { return it }
        }

        // 8. For generic elements with short text, use text content
        element.textContent?.trim()?.takeIf { it.isNotBlank() && it.length < 50 }?.let { return it }

        return null
    }

    /**
     * Compute element state from attributes
     */
    private fun computeState(attributes: Map<String, String>): Map<String, Boolean> {
        val state = mutableMapOf<String, Boolean>()

        // Boolean attributes
        if (attributes.containsKey("disabled")) state["disabled"] = true
        if (attributes.containsKey("readonly")) state["readonly"] = true
        if (attributes.containsKey("required")) state["required"] = true
        if (attributes.containsKey("checked")) state["checked"] = true
        if (attributes.containsKey("selected")) state["selected"] = true
        if (attributes.containsKey("hidden")) state["hidden"] = true
        if (attributes.containsKey("open")) state["expanded"] = true

        // ARIA states
        attributes["aria-disabled"]?.let { state["disabled"] = it == "true" }
        attributes["aria-hidden"]?.let { state["hidden"] = it == "true" }
        attributes["aria-expanded"]?.let { state["expanded"] = it == "true" }
        attributes["aria-selected"]?.let { state["selected"] = it == "true" }
        attributes["aria-checked"]?.let { 
            state["checked"] = when (it) {
                "true" -> true
                "mixed" -> true  // Partially checked
                else -> false
            }
        }
        attributes["aria-pressed"]?.let { state["pressed"] = it == "true" }
        attributes["aria-invalid"]?.let { state["invalid"] = it == "true" }
        attributes["aria-busy"]?.let { state["busy"] = it == "true" }

        return state
    }

    /**
     * Compute value for form controls
     */
    private fun computeValue(element: DOMElement): String? {
        val tagName = element.tagName.lowercase()
        val attrs = element.attributes

        return when (tagName) {
            "input", "textarea" -> attrs["value"]?.takeIf { 
                attrs["type"] != "password" && it.length < 100 
            }
            "select" -> null // Would need to check selected option
            "progress", "meter" -> attrs["value"]
            else -> null
        }
    }

    /**
     * Convert AccessibilityTree to compact JSON format for LLM consumption
     */
    fun toCompactJson(node: AccessibilityNode, indent: Int = 0): String {
        return buildString {
            appendNode(node, this, indent)
        }
    }

    private fun appendNode(node: AccessibilityNode, sb: StringBuilder, indent: Int) {
        val prefix = "  ".repeat(indent)
        sb.append(prefix)
        sb.append("{ role: \"${node.role}\"")
        
        node.name?.let { sb.append(", name: \"${escapeJson(it)}\"") }
        node.value?.let { sb.append(", value: \"${escapeJson(it)}\"") }
        
        if (node.state.isNotEmpty()) {
            val stateStr = node.state.entries
                .filter { it.value }
                .joinToString(", ") { "\"${it.key}\": true" }
            sb.append(", state: { $stateStr }")
        }
        
        if (node.children.isNotEmpty()) {
            sb.append(",\n$prefix  children: [\n")
            node.children.forEachIndexed { index, child ->
                appendNode(child, sb, indent + 2)
                if (index < node.children.size - 1) sb.append(",")
                sb.append("\n")
            }
            sb.append("$prefix  ]")
        }
        
        sb.append(" }")
    }

    private fun escapeJson(str: String): String {
        return str
            .replace("\\", "\\\\")
            .replace("\"", "\\\"")
            .replace("\n", "\\n")
            .replace("\r", "\\r")
            .replace("\t", "\\t")
    }

    /**
     * Filter the accessibility tree to only include actionable elements
     */
    fun filterActionable(node: AccessibilityNode): AccessibilityNode? {
        val actionableRoles = setOf(
            "button", "link", "textbox", "checkbox", "radio", "combobox",
            "listbox", "menuitem", "option", "slider", "spinbutton",
            "switch", "tab", "searchbox", "menuitemcheckbox", "menuitemradio"
        )

        val isActionable = node.role in actionableRoles
        val actionableChildren = node.children.mapNotNull { filterActionable(it) }

        return when {
            isActionable -> node.copy(children = actionableChildren)
            actionableChildren.isNotEmpty() -> {
                if (actionableChildren.size == 1) {
                    actionableChildren.first()
                } else {
                    node.copy(children = actionableChildren)
                }
            }
            else -> null
        }
    }

    /**
     * Get statistics about the accessibility tree
     */
    fun getStatistics(node: AccessibilityNode): AccessibilityStats {
        var totalNodes = 0
        var actionableNodes = 0
        val roleCounts = mutableMapOf<String, Int>()

        fun traverse(n: AccessibilityNode) {
            totalNodes++
            roleCounts[n.role] = (roleCounts[n.role] ?: 0) + 1
            
            val actionableRoles = setOf(
                "button", "link", "textbox", "checkbox", "radio", "combobox",
                "listbox", "menuitem", "option", "slider", "spinbutton",
                "switch", "tab", "searchbox"
            )
            if (n.role in actionableRoles) actionableNodes++
            
            n.children.forEach { traverse(it) }
        }

        traverse(node)

        return AccessibilityStats(
            totalNodes = totalNodes,
            actionableNodes = actionableNodes,
            roleCounts = roleCounts
        )
    }
}

/**
 * Accessibility tree node representation
 *
 * This is a semantic representation of UI elements suitable for LLM consumption.
 * Each node contains:
 * - role: ARIA role indicating the semantic purpose
 * - name: Accessible name for identification
 * - state: Current interactive states
 * - value: Current value for form controls
 * - selector: CSS selector for element targeting
 */
@Serializable
data class AccessibilityNode(
    val role: String,
    val name: String? = null,
    val state: Map<String, Boolean> = emptyMap(),
    val value: String? = null,
    val selector: String,
    val children: List<AccessibilityNode> = emptyList(),
    val description: String? = null
) {
    /**
     * Get a concise representation for display
     */
    fun toDisplayString(): String {
        val states = state.entries.filter { it.value }.map { it.key }
        val stateStr = if (states.isNotEmpty()) " [${states.joinToString(", ")}]" else ""
        val nameStr = name?.let { " \"$it\"" } ?: ""
        return "$role$nameStr$stateStr"
    }
}

/**
 * Statistics about an accessibility tree
 */
@Serializable
data class AccessibilityStats(
    val totalNodes: Int,
    val actionableNodes: Int,
    val roleCounts: Map<String, Int>
)