#!/usr/bin/env python
# coding:utf-8

"""
This file defines the nodes and logic needed for the plugin system in meshroom.
A plugin is a collection of node(s) of any type with their rutime environnement setup file attached.
We use the term 'environment' to abstract a docker container or a conda/virtual environment.
"""

from enum import Enum
import json
import os, sys
import logging
import urllib
from distutils.dir_util import copy_tree, remove_tree
import subprocess 
import venv
import inspect

from meshroom.core import desc
from meshroom.core import pluginsNodesFolder, pluginsPipelinesFolder, pluginCatalogFile, defaultCacheFolder
from meshroom.core import meshroomFolder
from meshroom.core.graph import loadGraph
from meshroom.core import hashValue

#executables def
meshroomBinDir = os.path.abspath(os.path.join(meshroomFolder, "..", "bin"))
condaBin = "conda"
dockerBin = "docker"

class EnvType(Enum):
    """
    enum for the type of env used (by degree of encapsulation)
    """
    NONE = 0
    PIP = 1
    REZ = 2
    VENV = 10
    CONDA = 20
    DOCKER = 30

#NOTE: could add the concept of dependencies between plugins
class PluginParams():
    """"
    Class that holds parameters to install one plugin from a folder and optionally from a json structure
    """
    def __init__(self, pluginUrl, jsonData=None):
        #get the plugin name from folder
        self.pluginName = os.path.basename(pluginUrl)
        #default node and pipeline locations
        self.nodesFolder = os.path.join(pluginUrl, "meshroomNodes")
        self.pipelineFolder = os.path.join(pluginUrl, "meshroomPipelines")
        #overwrite is json is passed
        if jsonData is not None:
            self.pluginName = jsonData["pluginName"]
            #default node and pipeline locations
            self.nodesFolder = os.path.join(pluginUrl, jsonData["nodesFolder"])
            if "pipelineFolder" in jsonData.keys():
                self.pipelineFolder = os.path.join(pluginUrl, jsonData["pipelineFolder"])

def getEnvName(envContent):
    return "meshroom_plugin_"+hashValue(envContent)

def _dockerImageExists(image_name, tag='latest'): 
    """
    Check if the desired image:tag exists 
    """
    try: 
        result = subprocess.run( [dockerBin, 'images', image_name, '--format', '{{.Repository}}:{{.Tag}}'], 
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True )
        if result.returncode != 0: 
            return False 
        images = result.stdout.splitlines() 
        image_tag = f"{image_name}:{tag}" 
        return image_tag in images 
    except Exception as e: 
        return False 

def _cleanEnvVarsRez():
    """
    Used to unset all rez defined env that mess up with conda.
    """
    cmd = "unset python; unset PYTHONPATH; "
    return cmd

def _condaEnvExist(envName):
    """
    Checks if a specified env exists
    """
    cmd = condaBin+" list --name "+envName+" > /dev/null"
    return os.system(cmd) == 0    

def _formatPluginName(pluginName):
    """
    Replaces spaces for env naming
    """
    return pluginName.replace(" ", "_")

def getVenvExe(venvPath):
    """
    Returns the path for the python in a virtual env 
    """
    if not os.path.isdir(venvPath):
        raise ValueError("The specified path "+venvPath+" is not a directory")
    if sys.platform == "win32":
        executable = os.path.join(venvPath, 'Scripts', 'python.exe')
    else:
        executable = os.path.join(venvPath, 'bin', 'python')
    if not os.path.isfile(executable):
        raise FileNotFoundError(f"Python executable not found in the specified virtual environment: "+executable)
    return executable

def getVenvPath(envName):
    return os.path.join(defaultCacheFolder, envName)

def _venvExists(envName):
    """
    Check if the following virtual env exists
    """
    return os.path.isdir(getVenvPath(envName))

def getActiveRezPackages():
    """
    Returns a list containing the active explicit packages
    """
    packages = []
    if 'REZ_REQUEST' in os.environ:
        nondefpackages = [n.split("-")[0] for n in os.environ.get('REZ_REQUEST', '').split()]
        resolvedPackages = os.environ.get('REZ_RESOLVE', '').split()
        resolvedVersions = {}
        for r in resolvedPackages:
            if r.startswith('~'):
                continue
            v = r.split('-')
            resolvedVersions[v[0]] = v[1]
        packages = [p+"-"+resolvedVersions[p] for p in resolvedVersions.keys() if p in nondefpackages]
    return packages

def installPlugin(pluginUrl):
    """
    Install plugin from an url or local path.
    Regardless of the method, the content will be copied in the plugin folder of meshroom (which is added to the list of directory to load nodes from).
    There are two options :
        - having the following structure :
            - [plugin folder] (will be the plugin name)
                - meshroomNodes
                    - [code for your nodes] that contains relative path to a DockerFile|env.yaml|requirements.txt
                    - [...]
                - meshroomPipelines
                    - [your meshroom templates]
        - having a meshroomPlugin.json file at the root of the plugin folder
          With this solution, you may have several separate node folders.
    """
    logging.info("Installing plugin from "+pluginUrl)
    try:
        isLocal = True
        #if git repo, clone the repo in cache
        if urllib.parse.urlparse(pluginUrl).scheme in ('http', 'https','git'):
            os.chdir(defaultCacheFolder)
            os.system("git clone "+pluginUrl)
            pluginName = pluginUrl.split('.git')[0].split('/')[-1]
            pluginUrl = os.path.join(defaultCacheFolder, pluginName)
            isLocal = False
        #sanity check
        if not os.path.isdir(pluginUrl):
            ValueError("Invalid plugin path :"+pluginUrl)
        #by default only one plugin, and with default file hierachy
        pluginParamList=[PluginParams(pluginUrl)]
        #location of the json file if any
        paramFile=os.path.join(pluginUrl, "meshroomPlugin.json")
        #load json for custom install if any
        if os.path.isfile(paramFile):
            jsonData=json.load(open(paramFile,"r"))
            pluginParamList = [PluginParams(pluginUrl, jsonDataplugin) for jsonDataplugin in jsonData]
        #for each plugin, run the 'install'
        for pluginParam in pluginParamList:
            intallFolder = os.path.join(pluginsNodesFolder, _formatPluginName(pluginParam.pluginName))
            logging.info("Installing "+pluginParam.pluginName+" from "+pluginUrl+" in "+intallFolder)
            #check if folder valid
            if not os.path.isdir(pluginParam.nodesFolder):
                raise RuntimeError("Invalid node folder: "+pluginParam.nodesFolder)
            #check if already installed
            if os.path.isdir(intallFolder):
                logging.warn("Plugin already installed, will overwrite")
                if os.path.islink(intallFolder):
                    os.unlink(intallFolder)
                else:
                    remove_tree(intallFolder)
            #install via symlink if local, otherwise copy (usefull to develop)
            if isLocal:
                os.symlink(pluginParam.nodesFolder, intallFolder)
                if os.path.isdir(pluginParam.pipelineFolder):
                    pipelineFolderLink = os.path.join(pluginsPipelinesFolder, pluginParam.pluginName)
                    if os.path.exists(pipelineFolderLink):
                        logging.warn("Plugin already installed, will overwrite")
                        os.unlink(pipelineFolderLink)
                    os.symlink(pluginParam.pipelineFolder, pipelineFolderLink)
            else:
                copy_tree(pluginParam.nodesFolder, intallFolder)
                if os.path.isdir(pluginParam.pipelineFolder):
                    copy_tree(pluginParam.pipelineFolder, os.path.join(pluginsPipelinesFolder, pluginParam.pluginName))
        #remove repo if was cloned
        if not isLocal:
            os.removedirs(pluginUrl)
        #NOTE: could try to auto load the plugins to avoid restart and test files
    except Exception as ex:
        logging.error(ex)
        return False
        
    return True

def getCatalog():
    """
    Returns the plugin catalog
    """
    jsonData=json.load(open(pluginCatalogFile,"r"))
    return jsonData
        
def getInstalledPlugin():
    """
    Returns the list of installed plugins
    """
    installedPlugins = [os.path.join(pluginsNodesFolder, f) for f in os.listdir(pluginsNodesFolder)]
    return installedPlugins

def uninstallPlugin(pluginUrl):
    """
    Uninstall a plugin
    """
    #NOTE: could also remove the env files
    if not os.path.exists(pluginUrl):
        raise RuntimeError("Plugin "+pluginUrl+" is not installed")
    if os.path.islink(pluginUrl):
        os.unlink(pluginUrl)
    else:
        os.removedirs(pluginUrl) 

def isBuilt(nodeDesc):
    """
    Check if the env needs to be build for a specific nodesc.
    """
    if nodeDesc.envType in [EnvType.NONE, EnvType.REZ]:
        return True
    elif nodeDesc.envType == EnvType.PIP:
        #NOTE: could find way to check for installed packages instead of rebuilding all the time
        return False
    elif nodeDesc.envType == EnvType.VENV:
        return _venvExists(nodeDesc._envName)
    elif nodeDesc.envType == EnvType.CONDA:
        return _condaEnvExist(nodeDesc._envName)
    elif nodeDesc.envType == EnvType.DOCKER:
        return _dockerImageExists(nodeDesc._envName)

def build(nodeDesc):
    """
    Perform the needed steps to prepare the environement in which to run the node.
    """
    if not hasattr(nodeDesc, 'envFile'):
        raise RuntimeError("The nodedesc has no env file")
    returnValue = 0
    if nodeDesc.envType in [EnvType.NONE, EnvType.REZ]:
        pass
    elif nodeDesc.envType == EnvType.PIP:
        #install packages in the same python as meshroom
        logging.info("Installing packages from "+ nodeDesc.envFile)
        buildCommand = sys.executable+" -m pip install "+ nodeDesc.envFile
        logging.info("Building with "+buildCommand+" ...")
        returnValue = os.system(buildCommand)
        logging.info("Done")
    elif nodeDesc.envType == EnvType.VENV:
        #create venv in default cache folder
        envPath = getVenvPath(nodeDesc._envName)
        logging.info("Creating virtual env "+envPath+" from "+nodeDesc.envFile)
        venv.create(envPath, with_pip=True)
        logging.info("Installing dependencies")
        envExe = getVenvExe(envPath)
        returnValue = os.system(_cleanEnvVarsRez()+envExe+" -m pip install -r "+ nodeDesc.envFile)
        venvPythonLibFolder = os.path.join(os.path.dirname(envExe), '..', 'lib')
        venvPythonLibFolder = [os.path.join(venvPythonLibFolder, p) 
                                for p in os.listdir(venvPythonLibFolder) if p.startswith("python")][0]
        os.symlink(meshroomFolder,os.path.join(venvPythonLibFolder, 'site-packages', 'meshroom'))
        logging.info("Done")
    elif nodeDesc.envType == EnvType.CONDA:
        #build a conda env from a yaml file
        logging.info("Creating conda env "+nodeDesc._envName+" from "+nodeDesc.envFile)
        makeEnvCommand = (  _cleanEnvVarsRez()+condaBin+" config --set channel_priority strict ; "
                            +condaBin+"  env create -v -v  --name "+nodeDesc._envName
                            +" --file "+nodeDesc.envFile+" ")
        logging.info("Making conda env")
        logging.info(makeEnvCommand)
        returnValue = os.system(makeEnvCommand)
        #find path to env's folder and add symlink to meshroom
        condaPythonExecudable=subprocess.check_output(_cleanEnvVarsRez()+condaBin+" run -n "+nodeDesc._envName
                                                        +" python -c \"import sys; print(sys.executable)\"",   
                                                        shell=True).strip().decode('UTF-8')
        condaPythonLibFolder=os.path.join(os.path.dirname(condaPythonExecudable), '..', 'lib')
        condaPythonLibFolder=[os.path.join(condaPythonLibFolder, p) 
                                for p in os.listdir(condaPythonLibFolder) if p.startswith("python")][0]
        os.symlink(meshroomFolder,os.path.join(condaPythonLibFolder, 'meshroom'))
        logging.info("Done making conda env")
    elif nodeDesc.envType == EnvType.DOCKER:
        #build docker image 
        logging.info("Creating image "+nodeDesc._envName+" from "+ nodeDesc.envFile)
        buildCommand = dockerBin+" build -f "+nodeDesc.envFile+" -t "+nodeDesc._envName+" "+os.path.dirname(nodeDesc.envFile)
        logging.info("Building with "+buildCommand+" ...")
        returnValue = os.system(buildCommand)
        logging.info("Done")
    else:
        raise RuntimeError("Invalid env type")
    if returnValue != 0:
        raise RuntimeError("Something went wrong during build")

def getCommandLine(chunk):
    """
    Return the command line needed to enter the environment + meshroom_compute
    Will make meshroom available in the environment.
    """
    nodeDesc=chunk.node.nodeDesc
    if chunk.node.isParallelized:
        raise RuntimeError("Parallelisation not supported for plugin nodes")    
    if chunk.node.graph.filepath == "":
        raise RuntimeError("The project needs to be saved to use plugin nodes")
    saved_graph = loadGraph(chunk.node.graph.filepath)
    if (str(chunk.node) not in [str(f) for f in saved_graph._nodes._objects]
        or chunk.node._uid != saved_graph.findNode(str(chunk.node))._uid ):
        raise RuntimeError("The changes needs to be saved to use plugin nodes")

    cmdPrefix = ""
    # vars common to venv and conda, that will be passed when runing conda run or venv
    meshroomCompute= meshroomBinDir+"/meshroom_compute"
    meshroomComputeArgs = "--node "+chunk.node.name+" \""+chunk.node.graph.filepath+"\""
    pythonsetMeshroomPath = "export PYTHONPATH="+meshroomFolder+":$PYTHONPATH;"

    if nodeDesc.envType == EnvType.VENV:
        envPath = getVenvPath(nodeDesc._envName)
        envExe = getVenvExe(envPath)
        #make sure meshroom in in pythonpath and that we call the right python 
        cmdPrefix = _cleanEnvVarsRez()+pythonsetMeshroomPath+" "+envExe + " "+ meshroomCompute +" " 
    elif nodeDesc.envType == EnvType.REZ:
        cmdPrefix = "rez env "+" ".join(getActiveRezPackages())+" "+nodeDesc._envName+" -- "+ meshroomCompute +" " 
    elif nodeDesc.envType == EnvType.CONDA:
        #NOTE: system env vars are not passed to conda run, we installed it 'manually' before
        cmdPrefix = _cleanEnvVarsRez()+condaBin+"  run --cwd "+os.path.join(meshroomFolder, "..")\
                    +" --no-capture-output -n "+nodeDesc._envName+" "+" python "+meshroomCompute
    elif nodeDesc.envType == EnvType.DOCKER:
        #path to the selected plugin
        classFile=inspect.getfile(chunk.node.nodeDesc.__class__)
        pluginDir = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(classFile),"..")))
        #path to the project/cache
        projectDir=os.path.abspath(os.path.realpath(os.path.dirname(chunk.node.graph.filepath)))
        mountCommand = (' --mount type=bind,source="'+projectDir+'",target="'+projectDir+'"' #mount with same file hierarchy
                        +' --mount type=bind,source="'+pluginDir+'",target=/meshroomPlugin,readonly' #mount the plugin folder (because of symbolic link, not necesseraly physically in meshroom's folder)
                        +' --mount type=bind,source="'+meshroomFolder+'",target=/meshroomFolder/meshroom,readonly' 
                        +' --mount type=bind,source="'+meshroomBinDir+'",target=/meshroomBinDir,readonly')
        #adds meshroom's code(& the plugin actual path that can be different (because of the dymbolic link)) to the python path
        envCommand = " --env PYTHONPATH=/meshroomFolder --env MESHROOM_NODES_PATH=/meshroomPlugin"
        #adds the gpu arg if needed
        runtimeArg=""
        if chunk.node.nodeDesc.gpu != desc.Level.NONE:
            runtimeArg="--runtime=nvidia --gpus all"
        #compose cl 
        cmdPrefix = dockerBin+" run -it --rm "+runtimeArg+" "+mountCommand+" "+envCommand+" "+nodeDesc._envName +" \"python /meshroomBinDir/meshroom_compute "
        meshroomComputeArgs="--node "+chunk.node.name+" "+chunk.node.graph.filepath+"\""
    else:
        raise RuntimeError("NodeType not recognised")

    command=cmdPrefix+" "+meshroomComputeArgs
        
    return command

# you may use these to explicitly define Pluginnodes
class PluginNode(desc.Node):
    pass

class PluginCommandLineNode(desc.CommandLineNode):
    pass