package dev.hgh.jai.reference

import com.intellij.openapi.diagnostic.Logger
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.PsiManager
import com.intellij.psi.PsiReferenceService
import com.intellij.psi.util.PsiTreeUtil
import dev.hgh.jai.lexer.JaiTokenTypes
import dev.hgh.jai.module.JaiModuleResolver
import dev.hgh.jai.module.JaiModuleTarget
import dev.hgh.jai.psi.JaiBlock
import dev.hgh.jai.psi.JaiDeclName
import dev.hgh.jai.psi.JaiDeclaration
import dev.hgh.jai.psi.JaiDirectiveExpr
import dev.hgh.jai.psi.JaiForName
import dev.hgh.jai.psi.JaiForStatement
import dev.hgh.jai.psi.JaiInitializer
import dev.hgh.jai.psi.JaiLiteralExpr
import dev.hgh.jai.psi.JaiParamName
import dev.hgh.jai.psi.JaiProcLiteralExpr
import dev.hgh.jai.psi.JaiRefExpr
import dev.hgh.jai.settings.JaiProjectSettings

private const val COMPLETION_DEBUG_PROPERTY = "jai.completion.debug"
private val LOG = Logger.getInstance("dev.hgh.jai.reference.JaiSymbolResolver")
private val RECOVERABLE_IMPORT =
    Regex("""(?m)(#(?:import|load))((?:\s*,\s*[A-Za-z_][A-Za-z0-9_]*)*)\s*"((?:\\.|[^"\\])*)"""")
private val RECOVERABLE_PROCEDURE =
    Regex("""(?s)(?:^|[;}\n])\s*[A-Za-z_][A-Za-z0-9_]*\s*::\s*\([^{}]*\)\s*(?:->[^\{;}]*)?\{""")
private val RECOVERABLE_LOCAL_DECLARATION =
    Regex("""(?m)(?:^|[;{}])\s*([A-Za-z_][A-Za-z0-9_]*)\s*(?::=|:(?!=))""")

/** A declaration-like PSI element that can be inserted by identifier completion. */
internal data class JaiSymbolCandidate(
    val element: PsiElement,
    val name: String,
    val kind: String = "symbol",
)

/** Shared symbol lookup for references and symbol-aware completion. */
internal object JaiSymbolResolver {
    private data class ScopedSymbol(
        val element: PsiElement,
        val name: String,
        val scope: PsiElement?,
        val fallbackScope: PsiElement?,
        val visibleFrom: Int,
    )

    private data class ImportBinding(
        val file: PsiFile,
        val directive: String,
        val scope: PsiElement?,
        val visibleFrom: Int,
        val boundName: String?,
        val importsIntoScope: Boolean,
        val excludedNames: Set<String>,
    )

    fun resolve(reference: JaiRefExpr): PsiElement? {
        val file = reference.containingFile ?: return null
        return completionCandidates(file, reference)
            .firstOrNull { it.name == reference.text }
            ?.element
    }

    /** Returns visible local declarations followed by visible unqualified imports. */
    fun completionCandidates(
        file: PsiFile,
        position: PsiElement,
    ): List<JaiSymbolCandidate> {
        val candidates = linkedMapOf<String, JaiSymbolCandidate>()

        val localSymbols = localSymbols(file)
        val recoveredLocals = recoverLocalSymbols(file, position)
        localSymbols
            .filter { isVisible(it, position) && !isInside(it.element, position) }
            .sortedWith(
                compareByDescending<ScopedSymbol> { scopeDepth(it.scope ?: it.fallbackScope) }
                    .thenBy { it.visibleFrom }
                    .thenBy { it.element.textOffset },
            ).forEach { symbol ->
                candidates.putIfAbsent(
                    symbol.name,
                    JaiSymbolCandidate(symbol.element, symbol.name),
                )
            }

        recoveredLocals.forEach { symbol ->
            candidates.putIfAbsent(symbol.name, symbol)
        }
        importedCandidates(file, position).forEach { symbol ->
            candidates.putIfAbsent(symbol.name, symbol)
        }
        val result = candidates.values.toList()
        debug(
            "file=${file.virtualFile?.path ?: file.name} roots=${JaiProjectSettings.getInstance(file.project).rootPaths()} " +
                "locals=${localSymbols.joinToString {
                    "${it.name}@${it.element.textOffset}/scope=${it.scope?.textRange}/fallback=${it.fallbackScope?.textRange}"
                }} recovered=${recoveredLocals.joinToString { it.name }} " +
                "imports=${importBindings(file).joinToString { it.file.virtualFile?.path ?: it.file.name }} " +
                "candidates=${result.take(40).joinToString { it.name }}",
        )
        return result
    }

    /** Returns members of a module bound to [qualifier] in the visible lexical scope. */
    fun memberCompletionCandidates(
        file: PsiFile,
        position: PsiElement,
        qualifier: String,
    ): List<JaiSymbolCandidate> {
        val binding =
            importBindings(file)
                .filter { it.boundName == normalizeName(qualifier) && isVisible(it, position) }
                .sortedWith(
                    compareByDescending<ImportBinding> { scopeDepth(it.scope) }
                        .thenByDescending { it.visibleFrom },
                ).firstOrNull()
                ?: return emptyList()

        return topLevelDeclarations(binding.file)
            .asSequence()
            .filter { symbolName(it) !in binding.excludedNames }
            .mapNotNull { name ->
                symbolName(name)?.let { JaiSymbolCandidate(name, it, "member") }
            }.distinctBy { it.name }
            .toList()
    }

    private fun localSymbols(file: PsiFile): List<ScopedSymbol> {
        val declarations =
            PsiTreeUtil
                .findChildrenOfType(file, JaiDeclaration::class.java)
                .flatMap { declaration ->
                    declaration.declNames.declNameList.mapNotNull { name ->
                        symbolName(name)?.let {
                            ScopedSymbol(
                                element = name,
                                name = it,
                                scope = PsiTreeUtil.getParentOfType(name, JaiBlock::class.java),
                                fallbackScope = PsiTreeUtil.getParentOfType(name, JaiProcLiteralExpr::class.java),
                                visibleFrom = name.textOffset,
                            )
                        }
                    }
                }

        val parameters =
            PsiTreeUtil
                .findChildrenOfType(file, JaiParamName::class.java)
                .mapNotNull { name ->
                    val symbol = symbolName(name) ?: return@mapNotNull null
                    val procedure = PsiTreeUtil.getParentOfType(name, JaiProcLiteralExpr::class.java)
                    ScopedSymbol(
                        element = name,
                        name = symbol,
                        // Keep the procedure itself as the parameter scope. The body block is a
                        // nested scope, so a local declaration correctly shadows a parameter.
                        scope = procedure,
                        fallbackScope = null,
                        visibleFrom = procedure?.textOffset ?: name.textOffset,
                    )
                }

        val forNames =
            PsiTreeUtil
                .findChildrenOfType(file, JaiForName::class.java)
                .mapNotNull { name ->
                    val symbol = symbolName(name) ?: return@mapNotNull null
                    ScopedSymbol(
                        element = name,
                        name = symbol,
                        scope = PsiTreeUtil.getParentOfType(name, JaiForStatement::class.java),
                        fallbackScope = PsiTreeUtil.getParentOfType(name, JaiProcLiteralExpr::class.java),
                        visibleFrom = name.textOffset,
                    )
                }

        return declarations + parameters + forNames
    }

    private fun importedCandidates(
        file: PsiFile,
        position: PsiElement,
    ): List<JaiSymbolCandidate> =
        importBindings(file)
            .asSequence()
            .filter { (it.boundName == null || it.importsIntoScope) && isVisible(it, position) }
            .flatMap { binding ->
                topLevelDeclarations(binding.file)
                    .asSequence()
                    .filter { symbolName(it) !in binding.excludedNames }
                    .mapNotNull { name ->
                        symbolName(name)?.let { JaiSymbolCandidate(name, it, "import") }
                    }
            }.distinctBy { it.name }
            .toList()

    /** Recovers simple locals when an unfinished statement prevents the procedure PSI from closing. */
    private fun recoverLocalSymbols(
        file: PsiFile,
        position: PsiElement,
    ): List<JaiSymbolCandidate> {
        val caretOffset = position.textOffset.coerceIn(0, file.textLength)
        val beforeCaret = file.text.substring(0, caretOffset)
        val procedure = RECOVERABLE_PROCEDURE.findAll(beforeCaret).lastOrNull() ?: return emptyList()
        val bodyStart = procedure.range.last + 1
        if (bodyStart >= beforeCaret.length) return emptyList()

        return RECOVERABLE_LOCAL_DECLARATION
            .findAll(beforeCaret.substring(bodyStart))
            .map { match -> JaiSymbolCandidate(position, match.groupValues[1], "local") }
            .distinctBy { it.name }
            .toList()
    }

    private fun importBindings(file: PsiFile): List<ImportBinding> {
        val parsed =
            PsiTreeUtil
                .findChildrenOfType(file, JaiLiteralExpr::class.java)
                .mapNotNull { literal ->
                    val directive = literal.parent as? JaiDirectiveExpr ?: return@mapNotNull null
                    val name = directive.directiveName() ?: return@mapNotNull null
                    if (name != "#import" && name != "#load") return@mapNotNull null

                    val declaration = boundDeclaration(literal)
                    val boundName =
                        declaration
                            ?.declNames
                            ?.declNameList
                            ?.firstOrNull()
                            ?.let(::symbolName)
                    val declarationPrefix = declaration?.let { prefixBeforeNames(it) }.orEmpty()
                    val importsIntoScope = declaration == null || declarationPrefix.trimStart().startsWith("using")
                    val excludedNames = excludedNames(declarationPrefix)
                    val target =
                        PsiReferenceService
                            .getService()
                            .getContributedReferences(literal)
                            .firstOrNull()
                            ?.resolve() as? PsiFile
                            ?: return@mapNotNull null

                    ImportBinding(
                        file = target,
                        directive = name,
                        scope = PsiTreeUtil.getParentOfType(literal, JaiBlock::class.java),
                        visibleFrom = literal.textOffset,
                        boundName = boundName,
                        importsIntoScope = importsIntoScope,
                        excludedNames = excludedNames,
                    )
                }

        // An unfinished statement immediately before the caret can make Grammar-Kit retain the
        // file's top-level declarations but omit later directive PSI. Recover bare imports from
        // the source text so completion still works while the user is typing.
        val recovered = if (parsed.isEmpty()) recoverImportBindings(file) else emptyList()
        return (parsed + recovered).distinctBy { binding ->
            listOf(
                binding.file.virtualFile?.path ?: binding.file.name,
                binding.scope?.textRange?.startOffset,
                binding.boundName,
                binding.visibleFrom,
            )
        }
    }

    private fun recoverImportBindings(file: PsiFile): List<ImportBinding> {
        val sourceFile = file.virtualFile ?: return emptyList()
        return RECOVERABLE_IMPORT
            .findAll(file.text)
            .mapNotNull { match ->
                val directive = match.groupValues[1]
                val flags =
                    Regex("""[A-Za-z_][A-Za-z0-9_]*""")
                        .findAll(match.groupValues[2])
                        .map { it.value }
                        .toSet()
                if ("string" in flags) return@mapNotNull null
                val path = unescapeImportPath(match.groupValues[3])
                val target =
                    JaiModuleResolver.resolve(
                        file.project,
                        sourceFile,
                        JaiModuleTarget(directive, path, flags),
                    ) ?: return@mapNotNull null
                val targetPsi = PsiManager.getInstance(file.project).findFile(target) ?: return@mapNotNull null
                ImportBinding(
                    file = targetPsi,
                    directive = directive,
                    scope = null,
                    visibleFrom = match.range.first,
                    boundName = null,
                    importsIntoScope = true,
                    excludedNames = emptySet(),
                )
            }.toList()
    }

    private fun unescapeImportPath(text: String): String =
        buildString(text.length) {
            var index = 0
            while (index < text.length) {
                if (text[index] == '\\' && index + 1 < text.length) {
                    append(text[index + 1])
                    index += 2
                } else {
                    append(text[index++])
                }
            }
        }

    /** The declaration that binds an import, not an enclosing procedure declaration. */
    private fun boundDeclaration(literal: JaiLiteralExpr): JaiDeclaration? {
        val initializer = PsiTreeUtil.getParentOfType(literal, JaiInitializer::class.java) ?: return null
        return PsiTreeUtil.getParentOfType(initializer, JaiDeclaration::class.java)
    }

    private fun prefixBeforeNames(declaration: JaiDeclaration): String {
        val start = declaration.textRange.startOffset
        val end = declaration.declNames.textRange.startOffset
        return declaration.containingFile.text.substring(start, end)
    }

    private fun excludedNames(prefix: String): Set<String> {
        val body = Regex("""except\s*\(([^)]*)\)""").find(prefix)?.groupValues?.get(1) ?: return emptySet()
        return body
            .split(',')
            .map { normalizeName(it.trim()) }
            .filter(String::isNotEmpty)
            .toSet()
    }

    /**
     * Returns declarations exported by a file, including files it textually loads. A module's
     * public surface is often assembled by loading several sibling files (for example, Basic
     * loads Print.jai). Do not follow #import here: an imported module's private dependencies
     * must not become unqualified candidates in the importing file.
     */
    private fun topLevelDeclarations(file: PsiFile): List<JaiDeclName> = topLevelDeclarations(file, linkedSetOf())

    private fun topLevelDeclarations(
        file: PsiFile,
        visited: MutableSet<String>,
    ): List<JaiDeclName> {
        val fileKey = file.virtualFile?.path ?: file.name
        if (!visited.add(fileKey)) return emptyList()

        val directDeclarations =
            PsiTreeUtil
                .findChildrenOfType(file, JaiDeclaration::class.java)
                .asSequence()
                .flatMap { declaration -> declaration.declNames.declNameList.asSequence() }
                .filter { PsiTreeUtil.getParentOfType(it, JaiBlock::class.java) == null }
                .sortedBy { it.textOffset }
                .toList()

        val loadedDeclarations =
            importBindings(file)
                .asSequence()
                .filter { it.directive == "#load" }
                .flatMap { topLevelDeclarations(it.file, visited).asSequence() }
                .toList()

        return directDeclarations + loadedDeclarations
    }

    private fun isVisible(
        symbol: ScopedSymbol,
        position: PsiElement,
    ): Boolean {
        val inScope =
            when {
                symbol.scope == null && symbol.fallbackScope == null -> true
                symbol.scope?.textRange?.contains(position.textOffset) == true -> true
                symbol.fallbackScope?.textRange?.contains(position.textOffset) == true -> true
                else -> false
            }
        if (!inScope) return false
        // Block-local declarations and loop names are not visible before their declaration. Keep
        // top-level declarations available for forward references, which Jai permits.
        if (symbol.scope != null && symbol.visibleFrom > position.textOffset) return false
        return true
    }

    private fun isVisible(
        binding: ImportBinding,
        position: PsiElement,
    ): Boolean =
        // File-scope imports are declarations for the whole file even when the directive is
        // written below its first use (a common Jai layout). Block-local imports remain ordered.
        (binding.scope == null || binding.scope.textRange.contains(position.textOffset)) &&
            (binding.scope == null || binding.visibleFrom <= position.textOffset)

    private fun isInside(
        element: PsiElement,
        position: PsiElement,
    ): Boolean = element == position || element.textRange.contains(position.textOffset)

    private fun scopeDepth(scope: PsiElement?): Int {
        var current = scope
        var depth = 0
        while (current != null) {
            if (current is JaiBlock || current is JaiForStatement || current is JaiProcLiteralExpr) depth++
            current = current.parent
        }
        return depth
    }

    private fun debug(message: String) {
        if (java.lang.Boolean.getBoolean(COMPLETION_DEBUG_PROPERTY)) LOG.info(message)
    }

    private fun symbolName(element: PsiElement): String? {
        val text = element.text.trim()
        if (text.isEmpty()) return null
        val name =
            when (element) {
                is JaiParamName -> text.removePrefix("$$").removePrefix("$")
                else -> text
            }
        return normalizeName(name).takeIf(String::isNotEmpty)
    }

    private fun normalizeName(name: String): String {
        val trimmed = name.trim()
        return if (trimmed.length >= 2 && trimmed.first() == '`' && trimmed.last() == '`') {
            trimmed.substring(1, trimmed.length - 1)
        } else {
            trimmed
        }
    }
}

private fun JaiDirectiveExpr.directiveName(): String? = node.findChildByType(JaiTokenTypes.DIRECTIVE)?.text