#!/usr/bin/env python
# Copyright 2019-2022 DADoES, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License in the root directory in the "LICENSE" file or at:
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import os
import subprocess
import anatools
import sys
from anatools.lib.channel import Channel, find_channelfile

parser = argparse.ArgumentParser(
    description="""
Use ana to run simulations with a Rendered.ai channel and graph.
    run with base settings:         ana
    run with specific channel:      ana --channel channelName
    run with specific graph:        ana --graph /path/to/graph.yaml
    run in preview mode:            ana --preview
""",
    formatter_class=argparse.RawDescriptionHelpFormatter
)
parser.add_argument('--channel', default=None)
parser.add_argument('--graph')
parser.add_argument('--loglevel', default="ERROR")
parser.add_argument('--logfile', default=None)
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--interp_num', type=int, default=0)
parser.add_argument('--batch_size', type=int, default=1)
parser.add_argument('--preview', action="store_true")
parser.add_argument('--output', type=str, default='./output')
parser.add_argument('--data', type=str, default='./data')
parser.add_argument('--version', action='store_true')
args = parser.parse_args()
if args.version: print(f'ana {anatools.__version__}'); sys.exit(1)

# Configure initial logging. Needed so errors in Channel class are displayed. Logging level
# may get overridden later
Channel.configure_logging(loglevel=args.loglevel, logfile=args.logfile, logfile_mode="w")

# get channel file name
if args.channel is None: args.channel = find_channelfile()
if args.channel is None: raise Exception('No channel file was specified or found.')
channel_file = args.channel
if os.path.splitext(channel_file)[1] == '': 
    if os.path.exists(channel_file + ".yml"): channel_file = channel_file + ".yml"
    elif os.path.exists(channel_file + ".yaml"): channel_file = channel_file + ".yaml"
    else: print('No channel file was specified or found.'); sys.exit(1)

channel = Channel(channel_file)
channel.process_args(args)

status = []
for i in range(args.batch_size):
    print(f'\nANA: Processing run {channel.execution_flags["--interp_num"]+i}')
    if channel.type == "blender":
        command = (
            f'blender --background --python {channel.ana_package_dir}/lib/blender_main.py -- \\\n' +
            f'--channel {channel_file} \\\n' +
            f'--graph {channel.execution_flags["--graph"]} \\\n' +
            f'--loglevel {channel.execution_flags["--loglevel"]} \\\n' +
            f'--interp_num {channel.execution_flags["--interp_num"]+i} \\\n' +
            f'--output {channel.execution_flags["--output"]} \\\n' +
            f'--data {channel.execution_flags["--data"]}'
        )
        if channel.execution_flags["--seed"] is not None:
            command = command + f' \\\n--seed {channel.execution_flags["--seed"]}'
        if channel.execution_flags["--preview"]:
            command = command + ' \\\n--preview'
        if channel.execution_flags["--logfile"] is not None:
            command = command + f' \\\n--logfile {channel.execution_flags["--logfile"]}'
    elif channel.type == "python":
        command = (
            f'python {channel.ana_package_dir}/lib/python_main.py \\\n' +
            f'--channel {channel.execution_flags["--channel"]} \\\n' +
            f'--graph {channel.execution_flags["--graph"]} \\\n' +
            f'--loglevel {channel.execution_flags["--loglevel"]} \\\n' +
            f'--interp_num {channel.execution_flags["--interp_num"]+i} \\\n' +
            f'--output {channel.execution_flags["--output"]} \\\n' +
            f'--data {channel.execution_flags["--data"]}'
        )
        if channel.execution_flags["--seed"] is not None:
            command = command + f' \\\n--seed {channel.execution_flags["--seed"]}'
        if channel.execution_flags["--preview"]:
            command = command + ' \\\n--preview'
        if channel.execution_flags["--logfile"] is not None:
            command = command + f' \\\n--logfile {channel.execution_flags["--logfile"]}'
    elif channel.type == "omniverse":
        command = (
            f'/startup.sh --ext-folder kit-exts/exts --enable ana.interpreter \\\n' +
            f'--/channel={channel.execution_flags["--channel"]} \\\n' +
            f'--/graph={channel.execution_flags["--graph"]} \\\n' +
            f'--/interp_num={channel.execution_flags["--interp_num"]+i} \\\n' +
            f'--/output={channel.execution_flags["--output"]} \\\n' +
            f'--/loglevel={channel.execution_flags["--loglevel"]} \\\n' +
            f'--/data={channel.execution_flags["--data"]}'
        )
        if channel.execution_flags["--seed"] is not None:
            command = command + f' \\\n--/seed={channel.execution_flags["--seed"]}'
        if channel.execution_flags["--preview"]:
            command = command + f' \\\n--/preview={channel.execution_flags["--preview"]}'
        if channel.execution_flags["--logfile"] is not None:
            command = command + f' \\\n--/logfile={channel.execution_flags["--logfile"]}'

    p = subprocess.run(command, shell=True)
    status.append(p.returncode)
for s in status:
    if s != 0: exit(s)
exit(0)