package dev.hgh.jai.completion

import com.intellij.codeInsight.completion.CompletionContributor
import com.intellij.codeInsight.completion.CompletionParameters
import com.intellij.codeInsight.completion.CompletionProvider
import com.intellij.codeInsight.completion.CompletionResultSet
import com.intellij.codeInsight.completion.CompletionType
import com.intellij.codeInsight.lookup.LookupElementBuilder
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.project.DumbAware
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.patterns.PlatformPatterns
import com.intellij.psi.tree.IElementType
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.util.ProcessingContext
import dev.hgh.jai.lexer.JaiLexer
import dev.hgh.jai.lexer.JaiTokenTypes
import dev.hgh.jai.module.JaiModuleResolver
import dev.hgh.jai.module.JaiPathMode
import dev.hgh.jai.psi.JaiDirectiveExpr
import dev.hgh.jai.reference.JaiSymbolResolver

/** Directives and their adjacent comma flags before an import/load string. */
private const val COMPLETION_DEBUG_PROPERTY = "jai.completion.debug"
private val LOG = Logger.getInstance("dev.hgh.jai.completion")
private val DIRECTIVE_WITH_FLAGS =
    Regex("""#(import|load)((?:\s*,\s*[A-Za-z_][A-Za-z0-9_]*)*)\s*$""")

/** Basic completion for Jai's context-free language and module names. */
class JaiCompletionContributor :
    CompletionContributor(),
    DumbAware {
    init {
        extend(
            CompletionType.BASIC,
            PlatformPatterns.psiElement(),
            JaiCompletionProvider(),
        )
    }
}

private class JaiCompletionProvider : CompletionProvider<CompletionParameters>() {
    override fun addCompletions(
        parameters: CompletionParameters,
        context: ProcessingContext,
        resultSet: CompletionResultSet,
    ) {
        val completionContext = JaiCompletionContext(parameters)
        val moduleDirective = completionContext.moduleDirective
        debug(
            "request file=${parameters.originalFile.virtualFile?.path ?: parameters.originalFile.name} " +
                "offset=${parameters.offset} token=${completionContext.debugToken} " +
                "prefix='${completionContext.prefix}' module=$moduleDirective " +
                "member=${completionContext.memberContext}",
        )
        when {
            moduleDirective != null -> {
                addModuleNames(parameters, moduleDirective, completionContext.prefix, resultSet)
            }

            completionContext.isCommentOrStringOrHereString -> {
                return
            }

            completionContext.memberContext != null -> {
                val member = completionContext.memberContext!!
                addMemberSymbols(parameters, member, resultSet)
            }

            completionContext.isDirectivePrefix -> {
                addDirectives(completionContext.prefix, resultSet)
            }

            else -> {
                addKeywords(completionContext.prefix, resultSet)
                addSymbols(parameters, completionContext.prefix, resultSet)
            }
        }
    }

    private fun addKeywords(
        prefix: String,
        resultSet: CompletionResultSet,
    ) {
        JaiTokenTypes.KEYWORD_MAP.keys
            .filter { it.startsWith(prefix) }
            .sorted()
            .forEach { keyword ->
                resultSet.addElement(
                    LookupElementBuilder
                        .create(keyword)
                        .withTypeText("keyword"),
                )
            }
    }

    private fun addDirectives(
        prefix: String,
        resultSet: CompletionResultSet,
    ) {
        JaiCompletionCatalog.DIRECTIVES
            .map { "#$it" }
            .filter { it.startsWith(prefix) }
            .forEach { directive ->
                val name = directive.removePrefix("#")
                resultSet.addElement(
                    LookupElementBuilder
                        // Keep the insertion text free of `#`: the prefix starts after
                        // `#` for a directive token, so inserting `#import` would yield
                        // `##import`. The alternate lookup string still matches `#im`.
                        .create(name)
                        .withPresentableText(directive)
                        .withLookupString(directive)
                        .withTypeText("directive"),
                )
            }
    }

    private fun addSymbols(
        parameters: CompletionParameters,
        prefix: String,
        resultSet: CompletionResultSet,
    ) {
        val candidates = JaiSymbolResolver.completionCandidates(parameters.originalFile, parameters.position)
        debug("unqualified candidates=${candidates.size} names=${candidates.take(40).joinToString { it.name }}")
        candidates
            .filter { it.name.startsWith(prefix) }
            .forEach { candidate ->
                resultSet.addElement(
                    LookupElementBuilder
                        .create(candidate.name)
                        .withPsiElement(candidate.element)
                        .withTypeText(candidate.kind),
                )
            }
    }

    private fun addMemberSymbols(
        parameters: CompletionParameters,
        member: JaiCompletionContext.MemberContext,
        resultSet: CompletionResultSet,
    ) {
        val candidates =
            JaiSymbolResolver.memberCompletionCandidates(
                parameters.originalFile,
                parameters.position,
                member.qualifier,
            )
        debug("member qualifier=${member.qualifier} candidates=${candidates.size} names=${candidates.take(40).joinToString { it.name }}")
        candidates
            .filter { it.name.startsWith(member.prefix) }
            .forEach { candidate ->
                resultSet.addElement(
                    LookupElementBuilder
                        .create(candidate.name)
                        .withPsiElement(candidate.element)
                        .withTypeText(candidate.kind),
                )
            }
    }

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

    private fun addModuleNames(
        parameters: CompletionParameters,
        directive: String,
        prefix: String,
        resultSet: CompletionResultSet,
    ) {
        val sourceFile = parameters.originalFile.virtualFile
        val project = parameters.originalFile.project
        val variants =
            if (directive == "#import") {
                JaiModuleCompletion.moduleNames(project, sourceFile)
            } else {
                JaiModuleCompletion.loadPaths(project, sourceFile)
            }
        variants
            .filter { it.startsWith(prefix) }
            .forEach { path ->
                resultSet.addElement(
                    LookupElementBuilder
                        .create(path)
                        .withTypeText(if (directive == "#import") "module" else "file"),
                )
            }
    }
}

private class JaiCompletionContext(
    private val parameters: CompletionParameters,
) {
    private val text =
        parameters.editor.document.charsSequence
            .toString()
    private val offset = parameters.offset.coerceIn(0, text.length)
    private val token = tokenAtCaret(text, offset)

    val isCommentOrStringOrHereString: Boolean
        get() =
            token?.type in
                setOf(
                    JaiTokenTypes.LINE_COMMENT,
                    JaiTokenTypes.BLOCK_COMMENT,
                    JaiTokenTypes.STRING,
                    JaiTokenTypes.HERE_STRING,
                    JaiTokenTypes.NOTE,
                )

    val isDirectivePrefix: Boolean
        get() = token?.type === JaiTokenTypes.DIRECTIVE || token?.type === JaiTokenTypes.HASH

    val debugToken: String
        get() = token?.let { "${it.type}@${it.start}..${it.end}" } ?: "none"

    data class MemberContext(
        val qualifier: String,
        val prefix: String,
    )

    val memberContext: MemberContext?
        get() =
            MEMBER_PREFIX.find(text.substring(0, offset))?.let { match ->
                MemberContext(
                    qualifier = match.groupValues[1],
                    prefix = match.groupValues.getOrNull(2).orEmpty(),
                )
            }

    val prefix: String
        get() {
            val current = token ?: return ""
            val start =
                when {
                    moduleDirective != null && current.type === JaiTokenTypes.STRING -> current.start + 1
                    isDirectivePrefix -> current.start
                    current.type === JaiTokenTypes.IDENT || current.type in JaiTokenTypes.KEYWORDS -> current.start
                    else -> return ""
                }
            return if (start <= offset) text.substring(start, offset) else ""
        }

    val moduleDirective: String?
        get() {
            if (token?.type !== JaiTokenTypes.STRING || token.end < offset) return null
            if (isAfterClosedString()) return null

            val beforeString = text.substring(0, token.start)
            val directiveMatch = DIRECTIVE_WITH_FLAGS.find(beforeString)
            val directive =
                PsiTreeUtil
                    .getParentOfType(parameters.position, JaiDirectiveExpr::class.java)
                    ?.node
                    ?.findChildByType(JaiTokenTypes.DIRECTIVE)
                    ?.text
                    ?: directiveMatch?.let { "#${it.groupValues[1]}" }
            if (directive != "#import" && directive != "#load") return null
            if (directive == "#import" && hasStringFlag(directiveMatch?.groupValues?.get(2).orEmpty())) {
                return null
            }
            return directive
        }

    private fun hasStringFlag(flags: String): Boolean =
        Regex("""[A-Za-z_][A-Za-z0-9_]*""")
            .findAll(flags)
            .any { it.value == "string" }

    private fun isAfterClosedString(): Boolean {
        if (token == null || token.end != offset || token.start >= offset) return false
        val tokenText = text.substring(token.start, offset)
        return tokenText.length > 1 && tokenText.last() == '"' && isUnescapedQuote(tokenText.lastIndex)
    }

    private fun isUnescapedQuote(index: Int): Boolean {
        var backslashes = 0
        var position = index - 1
        while (position >= 0 && text[token!!.start + position] == '\\') {
            backslashes++
            position--
        }
        return backslashes % 2 == 0
    }

    private data class LexedToken(
        val type: IElementType,
        val start: Int,
        val end: Int,
    )

    private companion object {
        val MEMBER_PREFIX =
            Regex("""(?:^|[^A-Za-z0-9_`])(`?[A-Za-z_][A-Za-z0-9_]*`?)\s*\.\s*(`?[A-Za-z_][A-Za-z0-9_]*`?)?$""")
    }

    private fun tokenAtCaret(
        text: String,
        offset: Int,
    ): LexedToken? {
        val lexer = JaiLexer()
        lexer.start(text, 0, offset, 0)
        var last: LexedToken? = null
        while (lexer.tokenType != null) {
            val current =
                LexedToken(
                    lexer.tokenType!!,
                    lexer.tokenStart,
                    lexer.tokenEnd,
                )
            last = current
            if (current.end >= offset) return current
            lexer.advance()
        }
        return last
    }
}

private object JaiCompletionCatalog {
    /** Core/compiler directives observed in the local Jai distribution. */
    val DIRECTIVES: List<String> =
        setOf(
            "add_context",
            "align",
            "as",
            "asm",
            "assert",
            "bake",
            "bake_arguments",
            "bake_constants",
            "body_text",
            "bytes",
            "c_call",
            "caller_location",
            "char",
            "code",
            "complete",
            "compiler",
            "cpp_method",
            "cpp_return_type_is_non_pod",
            "define",
            "deprecated",
            "discard",
            "dump",
            "else",
            "elsewhere",
            "expand",
            "file",
            "filepath",
            "foreign",
            "if",
            "ifdef",
            "import",
            "include",
            "insert",
            "intrinsic",
            "library",
            "load",
            "location",
            "module_parameters",
            "modify",
            "no_abc",
            "no_alias",
            "no_aoc",
            "no_context",
            "no_debug",
            "no_padding",
            "no_reset",
            "overlay",
            "place",
            "placeholder",
            "placeholders",
            "poke_name",
            "pragma",
            "procedure_of_call",
            "program_export",
            "run",
            "run_and_insert",
            "scope_export",
            "scope_file",
            "scope_module",
            "specified",
            "string",
            "symmetric",
            "this",
            "through",
            "type",
            "type_info_none",
            "type_info_no_size_complaint",
            "type_info_procedures_are_void_pointers",
            "undef",
            "version",
        ).sorted()
}

private object JaiModuleCompletion {
    private const val JAI_EXTENSION = "jai"

    fun moduleNames(
        project: com.intellij.openapi.project.Project,
        sourceFile: VirtualFile?,
    ): List<String> =
        collectCandidates(JaiModuleResolver.candidateRoots(project, sourceFile, JaiPathMode.MODULE)) { child ->
            when {
                child.isDirectory && child.findChild("module.jai") != null -> child.name
                !child.isDirectory && child.extension == JAI_EXTENSION -> child.nameWithoutExtension
                else -> null
            }
        }

    fun loadPaths(
        project: com.intellij.openapi.project.Project,
        sourceFile: VirtualFile?,
    ): List<String> =
        collectCandidates(JaiModuleResolver.candidateRoots(project, sourceFile, JaiPathMode.FILE)) { child ->
            if (!child.isDirectory && child.extension == JAI_EXTENSION) child.name else null
        }

    private fun collectCandidates(
        roots: List<VirtualFile>,
        nameOf: (VirtualFile) -> String?,
    ): List<String> {
        val names = linkedSetOf<String>()
        roots.forEach { root ->
            root.children
                .asSequence()
                .mapNotNull(nameOf)
                .forEach(names::add)
        }
        return names.sorted()
    }
}