# =============================================================================
#
#  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.
#
# =============================================================================
"""RefineMesh operation"""

import glob
import os
import stat
import string
import subprocess
import sys

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

OP_SUCCEEDED = int(smtk.operation.Operation.Outcome.SUCCEEDED)

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

# Make sure this folder is in sys.path
source_dir = os.path.dirname(__file__)
if source_dir not in sys.path:
    sys.path.insert(0, source_dir)

from foam_mixin import FoamMixin
from control_dict_writer import ControlDictWriter
from process_status import ProcessStatus


class RefineMesh(smtk.operation.Operation, FoamMixin):
    """Configures and optionally runs topoSet and refineMesh applications.

    Uses foam_operation.sbt for template.
    Uses the blockMesh case directory.
    """

    def __init__(self):
        smtk.operation.Operation.__init__(self)
        FoamMixin.__init__(self)
        # Do NOT store any smtk resources as member data (causes memory leak)
        self.workflows_folder = None
        self.case_dir = None

    def name(self):
        return "Setup and run topoSet and refineMesh"

    def createSpecification(self):
        spec = self._create_specification(app='refineMesh')
        return spec

    def operateInternal(self):
        """"""
        project = self._get_project()
        if project is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        att_resource = self._get_attribute_resource(project)
        if att_resource is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        self.workflows_folder = self._get_workflows_folder()
        if self.workflows_folder is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        # Check attributes
        att_names = ['controlDict', 'RefineMeshDirection']
        if not self._check_attributes(att_resource, att_names):
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        # Create case directory if needed
        project_dir = os.path.abspath(os.path.dirname(project.location()))
        case_dir = os.path.join(project_dir, 'foam/blockMesh')
        constant_dir = os.path.join(case_dir, 'constant')
        system_dir = os.path.join(case_dir, 'system')
        for path in [constant_dir, system_dir]:
            if not os.path.exists(path):
                os.makedirs(path)
        self.case_dir = case_dir

        # Generate controlDict file
        if not self.cd_writer.generate_controlDict(case_dir, att_resource):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Get region attributes
        att_list = att_resource.findAttributes('RefineMeshRegion')
        if not att_list:
            # Return "Fail" if no regions found
            self.log().addError("Error: no refinement regions specified")
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # The order that refinement is carried out is important.
        # For now, sort attributes by name
        att_list.sort(key=lambda att: att.name())

        # Get the topSetDictTemplate
        topoSet_template_file = os.path.join(
            self.workflows_folder, 'internal/foam/templates/topoSetDict')
        topoSet_string = None
        with open(topoSet_template_file) as ft:
            topoSet_string = ft.read()
        if topoSet_string is None:
            self.addError('Failed to load {}'.format(topoSet_template_file))

        # Generate topoSetDict files
        topoSet_template = string.Template(topoSet_string)
        for i, att in enumerate(att_list):
            box_item = att.findDouble('box')
            min_string = '({} {} {})'.format(
                box_item.value(0), box_item.value(2), box_item.value(4))
            max_string = '({} {} {})'.format(
                box_item.value(1), box_item.value(3), box_item.value(5))
            d = dict(min_coords=min_string, max_coords=max_string)
            topoSet_content = topoSet_template.substitute(d)

            filename = 'topoSetDict.{}'.format(i + 1)
            path = os.path.join(self.case_dir, 'system', filename)
            with open(path, 'w') as fw:
                fw.write(topoSet_content)

        # Generate refineMeshDict from template
        if not self._generate_refineMeshDict(self.case_dir, att_resource):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Generate Allrun & Allclean files
        if not self._write_allclean_file(self.case_dir):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        if not self._write_refinemesh_allrun_file(self.case_dir):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Check "Run" option
        run_item = self.parameters().findString('run')
        if run_item.isEnabled():
            if not self._run_foam(self.case_dir, run_item):
                return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        result = self._create_result()
        return result

    def _generate_refineMeshDict(self, case_dir: str, att_resource: smtk.attribute.Resource) -> bool:
        """Writes system/refineMeshDict file. Returns True on success."""
        # Load template file
        template_file = os.path.join(
            self.workflows_folder, 'internal/foam/templates/refineMeshDict')
        template_string = None
        with open(template_file) as ft:
            template_string = ft.read()
        if template_string is None:
            self.addError('Failed to load {}'.format(template_file))
            return False
        template = string.Template(template_string)

        # Generate input dictionary
        dir_att = att_resource.findAttribute('RefineMeshDirection')
        item_names = ['x', 'y', 'z']
        keywords = ['tan1', 'tan2', 'normal']
        keyword_list = list()
        for i, name in enumerate(item_names):
            dir_item = dir_att.find(name)
            if dir_item.isEnabled():
                keyword_list.append(keywords[i])

        # If nothing selected, use all 3 directions
        if not keyword_list:
            smtkInfoMessage(self.log(), 'No directions selected: using (1 1 1).')
            keyword_list = keywords.copy()
        directions = ' '.join(keyword_list)
        template_dict = dict(directions=directions)
        content = template.substitute(template_dict)

        complete = False
        system_dir = os.path.join(self.case_dir, 'system')
        refine_path = os.path.join(system_dir, 'refineMeshDict')
        with open(refine_path, 'w') as fp:
            fp.write(content)
            complete = True

        return complete

    def _write_refinemesh_allrun_file(self, case_dir):
        """Generate Allrun file for refineMesh"""
        wrote = False  # return value
        from foam_mixin import Allrun_CONTENT, CHECK_TEMPLATE

        # Get number of topoSetDict files
        pattern = case_dir + '/topoSetDict.*'
        topo_file_list = glob.glob(pattern)

        path = os.path.join(case_dir, 'Allrun')
        with open(path, 'w') as f:
            f.write(Allrun_CONTENT)
            f.write('runApplication blockMesh\n\n')

            if topo_file_list:
                range_end = len(topo_file_list) + 1
                int_list = range(1, range_end)
                str_list = [str(i) for i in int_list]
                str_range = ' '.join(str_list)

                f.write('for i in {}\n'.format(str_range))
                f.write('do\n')
                for i, topo_file in enumerate(topo_file_list):
                    f.write('    runApplication topoSet -dict system/topoSetDict.{}\n'.format(i + 1))
                    check_topo = CHECK_TEMPLATE.substitute(app='topoSet')
                    f.write(check_topo)

                    f.write('    runApplication refineMesh -dict system/refineMeshDict -overwrite\n')
                    check_refine = CHECK_TEMPLATE.substitute(app='refineMesh')
                    f.write(check_refine)

            f.write('done\n')
            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 _run_foam(self, case_dir, run_item):
        """Runs topoSet/refineSet sequence for each RefineMeshRegion attribute.

        HARD-CODED to only run synchronous
        """
        self.log().reset()
        smtk.InfoMessage(self.log(), 'Using case_dir {}'.format(case_dir))

        run_mode = run_item.value()
        if run_mode == 'async':  # launch and return
            self.log().addError('Error - async run_mode not supported')
            return False

        elif run_mode != 'sync':
            self.log().addError('internal error - unrecognized run_mode {}'.format(run_mode))
            return False

        args_prefix = []
        if hasattr(smtk, 'use_openfoam_docker') and smtk.use_openfoam_docker:
            args = self._docker_run_commands(case_dir)

        # Get the attribute resource
        ref_item = self.parameters().associations()
        project = ref_item.value()
        att_resource_set = project.resources().findByRole('attributes')
        att_resource = att_resource_set.pop()

        att_list = att_resource.findAttributes('RefineMeshRegion')
        if not att_list:
            self.log().addError('Internal error: calling _run_foam() with no RefineMeshRegion attributes')
            return False  # internal error

        system_dir = os.path.join(self.case_dir, 'system')
        refine_path = os.path.join(system_dir, 'refineMeshDict')
        if not os.path.exists(refine_path):
            self.log().addError('Internal error: file {} not found'.format(refine_path))

        log_dir = os.path.join(case_dir, 'logs')
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)

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

        for i in range(len(att_list)):
            # Check that topoSetDict file exists
            topo_filename = 'topoSetDict.{}'.format(i + 1)
            topo_path = os.path.join(system_dir, topo_filename)
            if not os.path.exists(topo_path):
                self.log().addError('Internal error: file {} not found'.format(topo_path))
                return False

            # Run topoSet
            logfilename = 'topoSet{}.log'.format(i + 1)
            logfile = os.path.join(log_dir, logfilename)
            dict_path = 'system/{}'.format(topo_filename)
            args = args_prefix + ['topoSet', '-dict', dict_path]
            with open(logfile, 'w') as fp:
                completed_proc = subprocess.run(
                    args, stdout=fp, stderr=fp, cwd=self.case_dir, env=env, universal_newlines=True)
                if completed_proc.returncode == 0:
                    self.status = ProcessStatus.Completed
                    success = True
                else:
                    self.log().addError('topoSet returned code {}'.format(completed_proc.returncode))
                    self.status = ProcessStatus.Error
                    return False

            # Run refineMesh -overwrite
            logfilename = 'refineMesh{}.log'.format(i + 1)
            logfile = os.path.join(log_dir, logfilename)
            args = args_prefix + ['refineMesh', '-overwrite']
            with open(logfile, 'w') as fp:
                completed_proc = subprocess.run(
                    args, stdout=fp, stderr=fp, cwd=self.case_dir, env=env, universal_newlines=True)
                if completed_proc.returncode == 0:
                    self.status = ProcessStatus.Completed
                    success = True
                else:
                    self.log().addError('refineMesh returned code {}'.format(completed_proc.returncode))
                    self.status = ProcessStatus.Error
                    return False

        return success
