#!/usr/bin/python3
import argparse
import logging
import os
import sys
from typing import NoReturn

try:
    import meshroom
except Exception:
    # If meshroom module is not in the PYTHONPATH, add our root using the relative path
    import pathlib
    meshroomRootFolder = pathlib.Path(__file__).parent.parent.resolve()
    sys.path.append(meshroomRootFolder)
    import meshroom
meshroom.setupEnvironment()

import meshroom.core
import meshroom.core.graph
from meshroom.core.node import Status
from meshroom.core.node import ChunkIndex


parser = argparse.ArgumentParser(description='Execute a Graph of processes.')
parser.add_argument('graphFile', metavar='GRAPHFILE.mg', type=str,
                    help='Filepath to a graph file.')
parser.add_argument('--node', metavar='NODE_NAME', type=str,
                    help='Process the node. It will generate an error if the dependencies are not already computed.')
parser.add_argument('--toNode', metavar='NODE_NAME', type=str,
                    help='Process the node with its dependencies.')
parser.add_argument('--inCurrentEnv', help='Execute process in current env without creating a dedicated runtime environment.',
                    action='store_true')
parser.add_argument('--forceStatus', help='Force computation if status is RUNNING or SUBMITTED.',
                    action='store_true')
parser.add_argument('--forceCompute', help='Compute in all cases even if already computed.',
                    action='store_true')
parser.add_argument('--extern', help='Use this option when you compute externally after submission to a render farm from meshroom.',
                    action='store_true')
parser.add_argument('--cache', metavar='FOLDER', type=str,
                    default=None,
                    help='Override the cache folder')
parser.add_argument('-v', '--verbose',
                    help='Set the verbosity level for logging:\n'
                            '  - fatal: Show only critical errors.\n'
                            '  - error: Show errors only.\n'
                            '  - warning: Show warnings and errors.\n'
                            '  - info: Show standard informational messages.\n'
                            '  - debug: Show detailed debug information.\n'
                            '  - trace: Show all messages, including trace-level details.',
                    default=os.environ.get('MESHROOM_VERBOSE', 'info'),
                    choices=['fatal', 'error', 'warning', 'info', 'debug', 'trace'])

parser.add_argument('-i', '--iteration', type=int, default=ChunkIndex.NONE, 
                    help='Define specific chunk index to compute')
parser.add_argument('--preprocess', help='Execute preprocess chunk', action='store_true')
parser.add_argument('--postprocess', help='Execute postprocess chunk', action='store_true')

args = parser.parse_args()

if args.preprocess:
    args.iteration = ChunkIndex.PREPROCESS
elif args.postprocess:
    args.iteration = ChunkIndex.POSTPROCESS

# Setup the verbose level
if args.extern:
    # For extern computation, we want to focus on the node computation log.
    # So, we avoid polluting the log with general warning about plugins, versions of nodes in file, etc.
    logging.getLogger().setLevel(level=logging.ERROR)
else:
    logging.getLogger().setLevel(meshroom.logStringToPython[args.verbose])

meshroom.core.initPlugins()
meshroom.core.initNodes()
meshroom.core.initSubmitters()

graph = meshroom.core.graph.loadGraph(args.graphFile)
if args.cache:
    graph.cacheDir = args.cache
graph.update()


def killRunningJob(node) -> NoReturn:
    """ Kills current job and try to avoid job restarting """
    jobInfo = node.nodeStatus.jobInfo
    submitterName = jobInfo.get("submitterName")
    if not submitterName:
        sys.exit(meshroom.MeshroomExitStatus.ERROR_NO_RETRY)
    from meshroom.core import submitters
    for subName, sub in submitters.items():
        if submitterName == subName:
            sub.killRunningJob()
            break
    sys.exit(meshroom.MeshroomExitStatus.ERROR_NO_RETRY)


if args.node:
    node = graph.findNode(args.node)
    node.updateStatusFromCache()
    submittedStatuses = [Status.RUNNING]

    if node.isCompatibilityNode:
        print(f'{node.name} is in Compatibility Mode and cannot be computed.')
        print(f'Compatibility issue: {node.issueDetails}')
        sys.exit(1)

    # Execute the node
    if not args.extern:
        # If running as "extern", the task is supposed to have the status SUBMITTED.
        # If not running as "extern", the SUBMITTED status should generate a warning.
        submittedStatuses.append(Status.SUBMITTED)

    if args.iteration >= 0 and not node._chunksCreated:
        # On a render farm, the nodeStatus file written by meshroom_createChunks may not be
        # immediately visible on this blade due to NFS propagation delays.
        # Wait for the file to appear and retry loading before giving up.
        import time
        retryInterval = 2    # seconds between retries
        maxRetryTime = 10    # total seconds before giving up
        elapsed = 0
        while not node._chunksCreated and elapsed < maxRetryTime:
            logging.warning(
                f"Chunks not created from cache for node {node.name} "
                f"(nodeStatusFile: \"{node.nodeStatusFile}\"). "
                f"Retrying in {retryInterval}s ({elapsed}/{maxRetryTime}s elapsed)..."
            )
            time.sleep(retryInterval)
            elapsed += retryInterval
            node.updateStatusFromCache()

        if not node._chunksCreated:
            logging.error(
                f"Computing chunk {args.iteration} of node {node} but chunks have not been "
                f"created after waiting {maxRetryTime}s. See file: \"{node.nodeStatusFile}\"."
            )
            sys.exit(-1)

    if node.isInitNode:
        print(f"InitNode: No computation to do.")
        sys.exit(0)

    if args.preprocess:
        chunks = [node.preprocessChunk]
    elif args.postprocess:
        chunks = [node.postprocessChunk]
    elif args.iteration == ChunkIndex.NONE:  # Default value
        chunks = node.chunks
    else:
        chunks = [node.chunks[args.iteration]]
    
    if not args.forceStatus and not args.forceCompute:
        for chunk in chunks:
            if chunk.status.status in submittedStatuses:
                # Particular case for the local isolated, the node status is set to RUNNING by the submitter directly.
                # We ensure that no other instance has started to compute, by checking that the computeSessionUid is empty.
                if chunk.node.getMrNodeType() == meshroom.core.MrNodeType.NODE and \
                    not chunk.status.computeSessionUid and node._nodeStatus.submitterSessionUid:
                    continue
                print(f'Warning: Node is already submitted with status "{chunk.status.status.name}". See file: "{chunk.statusFile}". ExecMode: {chunk.status.execMode.name}, computeSessionUid: {chunk.status.computeSessionUid}, submitterSessionUid: {node._nodeStatus.submitterSessionUid}')
                # sys.exit(-1)

    if args.extern:
        # Restore the log level
        logging.getLogger().setLevel(meshroom.logStringToPython[args.verbose])
        
    if args.iteration == ChunkIndex.NONE:
        # Process the whole node
        if node.nodeStatus.status == Status.STOPPED:
            print(f"Node {node}: status is STOPPED")
            killRunningJob(node)
        node.createChunks()
        node.prepareLogger(args.iteration)
        node.preprocess(args.forceCompute, args.inCurrentEnv)
        node.process(args.forceCompute, args.inCurrentEnv)
        node.postprocess(args.forceCompute, args.inCurrentEnv)
        node.restoreLogger()
    else:
        chunk = chunks[0]
        if chunk._status.status == Status.STOPPED:
            print(f"Chunk {chunk}: status is STOPPED")
            killRunningJob(node)
        node.prepareLogger(args.iteration)
        chunk.process(args.forceCompute, args.inCurrentEnv)
        node.restoreLogger()

else:
    if args.iteration != ChunkIndex.NONE:
        print('Error: "--iteration" only makes sense when used with "--node".')
        sys.exit(-1)
    toNodes = None
    if args.toNode:
        toNodes = graph.findNodes([args.toNode])

    meshroom.core.graph.executeGraph(graph, toNodes=toNodes, forceCompute=args.forceCompute, forceStatus=args.forceStatus)
