InputFileReference.kt

package nl.hannahsten.texifyidea.reference

import com.intellij.execution.RunManager
import com.intellij.execution.impl.RunManagerImpl
import com.intellij.openapi.roots.ProjectRootManager
import com.intellij.openapi.util.TextRange
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.PsiManager
import com.intellij.psi.PsiReferenceBase
import nl.hannahsten.texifyidea.algorithm.BFS
import nl.hannahsten.texifyidea.completion.pathcompletion.LatexGraphicsPathProvider
import nl.hannahsten.texifyidea.lang.commands.LatexGenericRegularCommand
import nl.hannahsten.texifyidea.psi.LatexCommands
import nl.hannahsten.texifyidea.psi.LatexPsiHelper
import nl.hannahsten.texifyidea.run.latex.LatexRunConfiguration
import nl.hannahsten.texifyidea.settings.sdk.LatexSdkUtil
import nl.hannahsten.texifyidea.util.LatexmkRcFileFinder
import nl.hannahsten.texifyidea.util.expandCommandsOnce
import nl.hannahsten.texifyidea.util.files.*
import nl.hannahsten.texifyidea.util.includedPackages
import nl.hannahsten.texifyidea.util.isTestProject
import nl.hannahsten.texifyidea.util.magic.CommandMagic

/**
 * Reference to a file, based on the command and the range of the filename within the command text.
 *
 * @param defaultExtension Default extension of the command in which this reference is.
 */
class InputFileReference(
    element: LatexCommands,
    val range: TextRange,
    val extensions: Set<String>,
    val defaultExtension: String
) : PsiReferenceBase<LatexCommands>(element) {

    init {
        rangeInElement = range
    }

    companion object {

        /**
         * Handle element rename, but taking into account whether the given
         * newElementName is just a filename which we have to replace,
         * or a full relative path (in which case we replace the whole path).
         */
        fun handleElementRename(command: LatexCommands, newElementName: String, elementNameIsJustFilename: Boolean): PsiElement {
            // A file has been renamed and we are given a new filename, to be replaced in the parameter text of the current command
            // It seems to be problematic to find the old filename we want to replace
            // Since the parameter content may be a path, but we are just given a filename, just replace the filename
            // We guess the filename is after the last occurrence of /
            val oldNode = command.node

            val newName = if ((oldNode?.psi as? LatexCommands)?.name in CommandMagic.illegalExtensions.keys) {
                newElementName.removeFileExtension()
            }
            else {
                newElementName
            }

            val defaultNewText = "${command.name}{$newName}"
            // Assumes that it is the last parameter, but at least leaves the options intact
            val default = oldNode?.text?.replaceAfterLast('{', "$newName}", defaultNewText) ?: defaultNewText

            // Recall that \ is a file separator on Windows
            val newText = if (elementNameIsJustFilename) {
                oldNode?.text?.trimStart('\\')?.replaceAfterLast('/', "$newName}", default.trimStart('\\'))
                    ?.let { "\\" + it } ?: default
            }
            else {
                default
            }
            val newNode = LatexPsiHelper(command.project).createFromText(newText).firstChild.node ?: return command
            if (oldNode == null) {
                command.parent?.node?.addChild(newNode)
            }
            else {
                command.parent.node.replaceChild(oldNode, newNode)
            }
            return command
        }
    }

    val key by lazy {
        rangeInElement.substring(element.text)
    }

    override fun resolve(): PsiFile? {
        return resolve(true)
    }

    /**
     * @param lookForInstalledPackages
     *              Whether to look for packages installed elsewhere on the filesystem.
     *              Set to false when it would make the operation too expensive, for example when trying to
     *              calculate the fileset of many files.
     * @param givenRootFile Used to avoid unnecessarily recalculating the root file.
     * @param isBuildingFileset
     *              True if we are building the fileset.
     *              If false we also need to resolve to graphics files. Doing so is really expensive at
     *              the moment (at least until the implementation in LatexGraphicsPathProvider is improved):
     *              for projects with 500 include commands in hundreds of files this can take 10 seconds in total if
     *              you call this function for every include command.
     *              However, note that doing only one resolve is not too expensive at all
     *              (10 seconds divided by 500 commands/resolves) so this is not a problem when doing only one resolve
     *              (if requested by the user).
     */
    fun resolve(lookForInstalledPackages: Boolean, givenRootFile: VirtualFile? = null, isBuildingFileset: Boolean = false): PsiFile? {
        // IMPORTANT In this method, do not use any functionality which makes use of the file set,
        // because this function is used to find the file set so that would cause an infinite loop

        // Get a list of extra paths to search in for the file, absolute or relative (to the directory containing the root file)
        val searchPaths = mutableListOf<String>()

        // Find the sources root of the current file.
        // findRootFile will also call getImportPaths, so that will be executed twice
        val rootFiles = if (givenRootFile != null) setOf(givenRootFile) else element.containingFile.findRootFiles().mapNotNull { it.virtualFile }
        val rootDirectories = rootFiles.mapNotNull { it.parent }

        // Check environment variables
        val runManager = RunManagerImpl.getInstanceImpl(element.project) as RunManager
        val texInputPath = runManager.allConfigurationsList
            .filterIsInstance<LatexRunConfiguration>()
            .firstOrNull { it.mainFile in rootFiles }
            ?.environmentVariables
            ?.envs
            ?.getOrDefault("TEXINPUTS", null)
            ?: LatexmkRcFileFinder.getTexinputsVariable(element.containingFile, null)

        if (texInputPath != null) {
            val path = texInputPath.trimEnd(':')
            searchPaths.add(path.trimEnd('/'))
            // See the kpathsea manual, // expands to subdirs
            if (path.endsWith("//")) {
                LocalFileSystem.getInstance().findFileByPath(path.trimEnd('/'))?.let { parent ->
                    searchPaths.addAll(
                        parent.allChildDirectories()
                            .filter { it.isDirectory }
                            .map { it.path }
                    )
                }
            }
        }

        // BIBINPUTS
        // Not used for building the fileset, so we can use the fileset to lookup the BIBINPUTS environment variable
        if (!isBuildingFileset && (element.name in CommandMagic.bibliographyIncludeCommands || extensions.contains("bib"))) {
            val bibRunConfigs = element.containingFile.getBibtexRunConfigurations()
            if (bibRunConfigs.any { config -> config.environmentVariables.envs.keys.any { it == "BIBINPUTS" } }) {
                // When using BIBINPUTS, the file will only be sought relative to BIBINPUTS
                searchPaths.clear()
                searchPaths.addAll(bibRunConfigs.mapNotNull { it.environmentVariables.envs["BIBINPUTS"] })
            }
        }

        var processedKey = expandCommandsOnce(key, element.project, file = rootFiles.firstOrNull()?.psiFile(element.project)) ?: key
        // Leading and trailing whitespaces seem to be ignored, at least it holds for \include-like commands
        processedKey = processedKey.trim()

        var targetFile: VirtualFile? = null

        // Try to find the target file directly from the given path
        @Suppress("KotlinConstantConditions")
        if (targetFile == null) {
            for (rootDirectory in rootDirectories) {
                targetFile = rootDirectory.findFile(filePath = processedKey, extensions = extensions, extensions.size != 6)
                if (targetFile != null) break
            }
        }

        // Try content roots
        if (targetFile == null && LatexSdkUtil.isMiktexAvailable) {
            for (moduleRoot in ProjectRootManager.getInstance(element.project).contentSourceRoots) {
                targetFile = moduleRoot.findFile(processedKey, extensions)
                if (targetFile != null) break
            }
        }

        // Try graphicspaths
        if (targetFile == null) {
            // If we are not building the fileset, we can make use of it
            if (!isBuildingFileset && element.containingFile.includedPackages().contains(LatexGenericRegularCommand.GRAPHICSPATH.dependency)) {
                // Add the graphics paths to the search paths
                searchPaths.addAll(LatexGraphicsPathProvider().getGraphicsPathsInFileSet(element.containingFile))
            }
            searchPath@ for (searchPath in searchPaths) {
                val path = if (!searchPath.endsWith("/")) "$searchPath/" else searchPath
                for (rootDirectory in rootDirectories) {
                    targetFile = rootDirectory.findFile(path + processedKey, extensions, false)
                    if (targetFile != null) break@searchPath
                }
            }
        }

        // Look for packages/files elsewhere using the kpsewhich command.
        if (targetFile == null && lookForInstalledPackages) {
            targetFile = element.getFileNameWithExtensions(processedKey)
                .mapNotNull { LatexPackageLocationCache.getPackageLocation(it, element.project) }
                .firstNotNullOfOrNull { getExternalFile(it) }
        }

        if (targetFile == null) targetFile = searchFileByImportPaths(element)?.virtualFile

        // \externaldocument uses the .aux file in the output directory, we are only interested in the source file, but it can be anywhere (because no relative path will be given, as in the output directory everything will be on the same level).
        // This does not count for building the file set, because the external document is not actually in the fileset, only the label definitions are
        if (!isBuildingFileset && targetFile == null && element.name == LatexGenericRegularCommand.EXTERNALDOCUMENT.commandWithSlash) {
            targetFile = findAnywhereInProject(processedKey)
        }

        if (targetFile == null) return null

        // Return a reference to the target file.
        return PsiManager.getInstance(element.project).findFile(targetFile)
    }

    /**
     * Try to find the file anywhere in the project. Returns the first match.
     * Might be expensive for large projects because of recursively visiting all directories, not sure.
     */
    fun findAnywhereInProject(fileName: String): VirtualFile? {
        val basePath = if (element.project.isTestProject().not()) {
            LocalFileSystem.getInstance().findFileByPath(element.project.basePath ?: return null) ?: return null
        }
        else {
            element.containingFile.virtualFile.parent ?: return null
        }
        BFS(basePath, { file -> file.children.toList() }).apply {
            iterationAction = { file: VirtualFile ->
                if (file.nameWithoutExtension == fileName && file.extension in extensions) {
                    BFS.BFSAction.ABORT
                }
                else {
                    BFS.BFSAction.CONTINUE
                }
            }
            execute()
            return end
        }
    }

    override fun handleElementRename(newElementName: String): PsiElement {
        return handleElementRename(element, newElementName, true)
    }

    // Required for moving referenced files
    override fun bindToElement(givenElement: PsiElement): PsiElement {
        val newFile = givenElement as? PsiFile ?: return this.element
        // Assume LaTeX will accept paths relative to the root file
        val newFileName = newFile.virtualFile?.path?.toRelativePath(this.element.containingFile.findRootFile().virtualFile.parent.path) ?: return this.element
        return handleElementRename(element, newFileName, false)
    }

    /**
     * Create a set possible complete file names (including extension), based on
     * the command that includes a file, and the name of the file.
     */
    private fun LatexCommands.getFileNameWithExtensions(fileName: String): Set<String> {
        val extension = CommandMagic.includeOnlyExtensions[this.commandToken.text] ?: emptySet()
        return extension.map { "$fileName.$it" }.toSet() + setOf(fileName)
    }
}