intellijai/src/main/kotlin/dev/hgh/jai/module/JaiModuleResolver.kt

package dev.hgh.jai.module

import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.vfs.VirtualFile
import dev.hgh.jai.settings.JaiProjectSettings

/** The information needed to resolve an #import or #load string. */
data class JaiModuleTarget(
    val directive: String,
    val path: String,
    val flags: Set<String>,
)

enum class JaiPathMode {
    MODULE,
    FILE,
    DIRECTORY,
    RELATIVE,
}

/** Resolves Jai paths and exposes the same root order to completion. */
object JaiModuleResolver {
    private val fileSystem: LocalFileSystem
        get() = LocalFileSystem.getInstance()

    fun resolve(
        project: Project,
        sourceFile: VirtualFile,
        target: JaiModuleTarget,
    ): VirtualFile? {
        val rawPath = target.path.replace('\\', '/')
        val path = normalize(rawPath)
        if (path.isEmpty()) return null

        val mode =
            when {
                target.directive == "#load" || "file" in target.flags -> JaiPathMode.FILE
                "dir" in target.flags -> JaiPathMode.DIRECTORY
                rawPath.startsWith("./") || rawPath.startsWith("../") -> JaiPathMode.RELATIVE
                else -> JaiPathMode.MODULE
            }
        if (path.startsWith('/')) {
            return findAbsolute(path, mode)
        }

        for (root in candidateRoots(project, sourceFile, mode)) {
            findInRoot(root, path, mode)?.let { return it }
        }
        return null
    }

    /**
     * Returns roots in resolution precedence order. Configured roots are intentionally shared by
     * the resolver and completion so a path cannot resolve differently from the path offered by
     * completion.
     */
    fun candidateRoots(
        project: Project,
        sourceFile: VirtualFile?,
        mode: JaiPathMode,
    ): List<VirtualFile> {
        val roots = linkedMapOf<String, VirtualFile>()

        fun add(root: VirtualFile?) {
            if (root != null && root.isValid && root.isDirectory) {
                roots.putIfAbsent(root.path, root)
            }
        }

        val sourceRoot = sourceFile?.parent
        when (mode) {
            JaiPathMode.MODULE -> add(sourceRoot?.findChild("modules"))
            JaiPathMode.FILE, JaiPathMode.DIRECTORY, JaiPathMode.RELATIVE -> add(sourceRoot)
        }

        // A configured root is an explicit opt-in to searching outside the project. For direct
        // #load/#import paths it is searched after the containing directory; for bare module
        // imports it wins over the built-in/project fallbacks below.
        if (mode != JaiPathMode.RELATIVE) {
            JaiProjectSettings.getInstance(project).configuredRoots().forEach(::add)
        }

        if (mode == JaiPathMode.MODULE) {
            val projectRoot = project.basePath?.let(fileSystem::findFileByPath)
            add(projectRoot?.findChild("modules"))
            add(projectRoot)
            add(
                fileSystem.findFileByPath(
                    "${System.getProperty("user.home")}/.local/jai/modules",
                ),
            )
        }
        return roots.values.toList()
    }

    private fun findAbsolute(
        path: String,
        mode: JaiPathMode,
    ): VirtualFile? {
        val direct = fileSystem.findFileByPath(path)
        if (direct != null) {
            when (mode) {
                JaiPathMode.FILE -> {
                    if (!direct.isDirectory) return direct
                }

                JaiPathMode.DIRECTORY -> {
                    if (direct.isDirectory) return direct.findChild("module.jai")
                }

                JaiPathMode.RELATIVE -> {
                    return if (direct.isDirectory) direct.findChild("module.jai") else direct
                }

                JaiPathMode.MODULE -> {
                    if (!direct.isDirectory) return direct
                    direct.findChild("module.jai")?.let { return it }
                }
            }
        }
        if (mode == JaiPathMode.MODULE && !path.endsWith(".jai")) {
            return fileSystem.findFileByPath("$path.jai")
                ?: fileSystem.findFileByPath("$path/module.jai")
        }
        return null
    }

    private fun findInRoot(
        root: VirtualFile,
        path: String,
        mode: JaiPathMode,
    ): VirtualFile? {
        val direct = root.findFileByRelativePath(path)
        if (direct != null) {
            when (mode) {
                JaiPathMode.FILE -> {
                    if (!direct.isDirectory) return direct
                }

                JaiPathMode.DIRECTORY -> {
                    if (direct.isDirectory) return direct.findChild("module.jai")
                }

                JaiPathMode.RELATIVE -> {
                    return if (direct.isDirectory) direct.findChild("module.jai") else direct
                }

                JaiPathMode.MODULE -> {
                    if (!direct.isDirectory) return direct
                    direct.findChild("module.jai")?.let { return it }
                }
            }
        }
        if (mode == JaiPathMode.MODULE && !path.endsWith(".jai")) {
            return root.findFileByRelativePath("$path.jai")
                ?: root.findFileByRelativePath("$path/module.jai")
        }
        return null
    }

    private fun normalize(path: String): String = path.replace('\\', '/').removePrefix("./")
}