# =============================================================================
#
#  Copyright (c) Kitware, Inc.
#  All rights reserved.
#  See LICENSE.txt for details.
#
#  This software is distributed WITHOUT ANY WARRANTY; without even
#  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
#  PURPOSE.  See the above copyright notice for more information.
#
# =============================================================================
"""Common utilities for openfoam operations"""

import os
import stat
import string
import subprocess
import sys
from typing import Union

import smtk
import smtk.attribute
import smtk.io
import smtk.operation
import smtk.project


# Make sure __file__ is set when using modelbuilder
import inspect
source_file = os.path.abspath(inspect.getfile(inspect.currentframe()))
__file__ = source_file

from control_dict_writer import ControlDictWriter
from process_status import ProcessStatus

OPENFOAM_VERSION = '2312'


# Commen text for Allclean files
Allclean_CONTENT = \
    """#!/bin/sh

cd "${0%/*}" || exit                            # Run from this directory
. ${WM_PROJECT_DIR:?}/bin/tools/CleanFunctions  # Tutorial clean functions

#--------------------------------------------------------------------------#

cleanCase0
rm -rf constant/polyMesh
"""

# Common text for Allrun files
Allrun_CONTENT = \
    """#!/bin/sh

cd "${0%/*}" || exit                            # Run from this directory
. ${WM_PROJECT_DIR:?}/bin/tools/RunFunctions    # Tutorial run functions

#--------------------------------------------------------------------------#
rm -f log.*

"""

check_template_content = """
if [ $$? -ne 0 ]; then
  echo "${app} Failed"
  exit 1
fi

"""
CHECK_TEMPLATE = string.Template(check_template_content)


class DictTemplate(string.Template):
    """String template with custom delimiter for operations.

    OpenFOAM uses $ to substitute variables from include files.
    This template uses a comment delimiter '//@' to insert strings
    generated by operations.
    """
    delimiter = '//@'


class FoamMixin:
    """Python MIXIN CLASS for operations that run OpenFoam applications.

    Provides a set of common methods used by many operations.
    Not implemented as an smtk.operation.Operation subclass because smtk cannot find
    operateInternal() method of its subclasses.
    """

    def __init__(self):
        # Do NOT store any smtk resources as member data (causes memory leak)
        self.source_dir = os.path.abspath(os.path.dirname(__file__))
        self.cd_writer = ControlDictWriter()  # for subclass

        # Operation results data
        self.case_dir = ''  # updated when operation runs
        self.status = 0
        self.pid = 0
        self.logfile = ''

    def _create_specification(self, app=None, sbt_file='foam_operation.sbt'):
        """Create default specification.

        Inputs
            app: if specified, update "Run" label to use app name
            sbt_file: template file to create specification
        """
        spec = self.createBaseSpecification()

        sbt_path = os.path.join(self.source_dir, sbt_file)
        if not os.path.exists(sbt_path):
            message = 'File not found: {}'.format(sbt_path)
            self.log().addError(message)
            raise RuntimeError(message)

        reader = smtk.io.AttributeReader()
        hasErr = reader.read(spec, sbt_path, self.log())
        if hasErr:
            message = 'Error loading specification file: {}'.format(sbt_path)
            self.log().addError(message)
            raise RuntimeError(message)

        # Update run label if specified
        if app is not None:
            att_defn = spec.findDefinition('setup')
            pos = att_defn.findItemPosition('run')
            if pos >= 0:
                run_defn = att_defn.itemDefinition(pos)
                label = 'Run {}'.format(app)
                run_defn.setLabel(label)
        return spec

    def _get_project(self) -> smtk.project.Project:
        """Gets project from parameter association, or returns None if not found."""
        err_message = None
        ref_item = self.parameters().associations()
        if ref_item is None:
            err_message = 'internal error: project association item is missing'
        elif ref_item.numberOfValues() != 1:
            err_message = 'should have 1 project association not {}'.ref_item.numberOfValues()
        if err_message:
            self.log().addError(err_message)
            return None

        project = ref_item.value()
        return project

    def _get_attribute_resource(self, project: smtk.project.Project) -> smtk.attribute.Resource:
        """Get project's attribute resource."""
        att_resource_set = project.resources().findByRole('attributes')
        if not att_resource_set:
            self.log().addError('Project missing attribute resource.')
            return None

        att_resource = att_resource_set.pop()
        return att_resource

    def _get_workflows_folder(self):
        """"""
        # Application must use SetupPaths op to assign smtk.workflows_folder
        if not hasattr(smtk, 'workflows_folder'):
            self.log().addError(
                'smtk is missing workflows_folder. Do you need to set DEVELOPER_MODE?')
            return None
        return smtk.workflows_folder

    def _check_attributes(self, att_resource, att_names: list) -> bool:
        """Check instanced attributes validity"""
        ok = True  # return value
        for att_name in att_names:
            att = att_resource.findAttribute(att_name)
            if att is None:
                self.log().addError('Did not find attribute {}'.format(att_name))
            elif not att.isValid():
                self.log().addError('Attribute {} is not valid'.format(att_name))
        return ok

    def _docker_run_commands(self, working_dir: str) -> list:
        """Generates command and args to launch openfoam container"""
        engine = 'docker'
        if hasattr(smtk, 'container_engine'):
            engine = smtk.container_engine
        args_list = [engine, 'run', '--rm']
        args_list.append('--user={}:{}'.format(smtk.uid, smtk.gid))
        args_list.append('--volume={}:/home/openfoam'.format(working_dir))
        args_list.append('opencfd/openfoam-run:{}'.format(OPENFOAM_VERSION))
        return args_list

    def _run_openfoam(self,
                      app: str,
                      case_dir: str,
                      run_item: smtk.attribute.IntItem,
                      args_list: list = [],
                      foamfile: str = None) -> bool:
        """Runs openfoam application from case directory."""
        self.case_dir = case_dir
        log_dir = os.path.join(case_dir, 'logs')
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
        self.logfile = os.path.join(log_dir, '{}.log'.format(app))
        errfile = os.path.join(log_dir, 'stderr.log')

        # Set PWD for standard (non-docker) OpenFOAM binaries
        env = os.environ.copy()
        env['PWD'] = case_dir

        args = [app] + args_list
        # Check hack for docker implementation
        if hasattr(smtk, 'use_openfoam_docker') and smtk.use_openfoam_docker:
            # args = ['openfoam-docker', '/'] + args
            args = self._docker_run_commands(case_dir) + args
            args_string = ' '.join(args)
            msg = 'Run command: {}'.format(args_string)
            self.log().addRecord(smtk.io.Logger.Info, msg)

        run_mode = run_item.value()
        if run_mode == 'sync':  # run and wait for completion
            with open(self.logfile, 'w') as fp:
                completed_proc = subprocess.run(
                    args, cwd=case_dir, env=env, universal_newlines=True)
                if completed_proc.returncode == 0:
                    self.status = ProcessStatus.Completed
                    success = True

                    if foamfile is not None:
                        if not foamfile.endswith('.foam'):
                            foamfile = '{}.foam'.format(foamfile)
                        foamfile_path = os.path.join(case_dir, foamfile)
                        with open(foamfile_path, 'a'):
                            os.utime(foamfile_path, None)

                else:
                    self.log().addError('{} returned code {}'.format(app, completed_proc.returncode))
                    self.status = ProcessStatus.Error
                    success = False

            return success

        elif run_mode == 'async':  # launch and return
            with open(self.logfile, 'w') as out, open(errfile, 'w') as err:
                proc = subprocess.Popen(
                    args, stdout=out, stderr=err, cwd=case_dir, env=env, universal_newlines=True)
                self.log().addRecord(smtk.io.Logger.Info, 'Started process {}'.format(proc.pid))
                self.status = ProcessStatus.Started
                self.pid = proc.pid
            return True

        else:
            self.log().addError('internal error - unrecognized run_mode {}'.format(run_mode))
            return False

    def _create_result(self):
        """Creates SUCCEEDED result and sets process-related items."""
        result = self.createResult(smtk.operation.Operation.Outcome.SUCCEEDED)
        result.findDirectory('case-directory').setValue(self.case_dir)
        result.findInt('status').setValue(self.status)
        result.findInt('pid').setValue(self.pid)
        result.findString('logfile').setValue(self.logfile)
        return result

    def _write_allclean_file(self, case_dir: str) -> bool:
        """Writes Allclean file to case directory"""
        wrote = False
        path = os.path.join(case_dir, 'Allclean')
        with open(path, 'w') as f:
            f.write(Allclean_CONTENT)
            wrote = True

        if not wrote:
            self.log().addError('Error writing {}'.format(path))
            return False

        # Set execute permissions
        st = os.stat(path)
        os.chmod(path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)

        return True

    def _write_allrun_file(
            self,
            case_dir: str,
            foam_apps: Union[str, list],
            copy_casename=None,
            restore0=False) -> bool:
        """Writes Allrun file to case directory

        Param foam_apps can be single app or list
        Param copy_casename if to first copy polyMesh from different case
        """
        if isinstance(foam_apps, str):
            foam_apps = [foam_apps]

        wrote = False
        path = os.path.join(case_dir, 'Allrun')
        with open(path, 'w') as f:
            f.write(Allrun_CONTENT)

            if copy_casename is not None:
                rm_cmd = 'rm -rf ./constant/polyMesh\n'
                f.write(rm_cmd)

                cp_cmd = 'cp -r ../{}/constant/polyMesh ./constant'.format(copy_casename)
                f.write(cp_cmd)
                check_lines = CHECK_TEMPLATE.substitute(app=cp_cmd)
                f.write(check_lines)

            if restore0:
                f.write('rm -rf 0\n')
                cp_cmd = 'cp -r 0.orig 0'
                f.write(cp_cmd)
                check_lines = CHECK_TEMPLATE.substitute(app=cp_cmd)
                f.write(check_lines)

            for app in foam_apps:
                f.write('runApplication {}'.format(app))
                check_lines = CHECK_TEMPLATE.substitute(app=app)
                f.write(check_lines)
            wrote = True

        if not wrote:
            self.log().addError('Error writing {}'.format(path))
            return False

        # Set execute permissions
        st = os.stat(path)
        os.chmod(path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)

        return True
