# =============================================================================
#
#  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.
#
# =============================================================================

"""icoFoam solver operation"""

import glob
import io
import os
import shutil
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, DictTemplate


class IcoFoam(smtk.operation.Operation, FoamMixin):
    """Generates inputs files for project analysis mesh and runs OpenFOAM solver."""

    def __init__(self):
        smtk.operation.Operation.__init__(self)
        FoamMixin.__init__(self)

        # Do NOT store any smtk resources as member data (causes memory leak)
        self.project_dir = None
        self.case_dir = None
        self.workflows_folder = None

    def name(self):
        return "Setup and run OpenFOAM solver"

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

    def operateInternal(self):
        """Generate the OpenFOAM input files for snappyHexMesh

        Steps include
          * Copy mesh file from snappyHexMesh
          * Copy 0.orig directory to 0
          * Generate p and U files from templates
          * Copy system files fvSchemes, fvSolution to system directory
          * Generate 0/include/initialConditions file
        """
        project = self._get_project()
        if project is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)
        self.project_dir = os.path.abspath(os.path.dirname(project.location()))

        att_resource = self._get_attribute_resource(project)
        # print('att_resource:', att_resource)
        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)

        solver = self.cd_writer.get_application_name(att_resource)
        self.case_dir = os.path.join(self.project_dir, 'foam', solver)

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

        self._reset_directory()

        # Copy mesh from snappyHexMesh case
        from_dir = os.path.join(self.case_dir, os.pardir, 'snappyHexMesh/constant/polyMesh')
        if not os.path.exists(from_dir):
            self.log().addError('No mesh at snappyMesh/constant/polyMesh')
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)
        to_dir = os.path.join(self.case_dir, 'constant/polyMesh')
        if os.path.exists(to_dir):
            shutil.rmtree(to_dir)
        shutil.copytree(from_dir, to_dir)

        self._copy_system_files(att_resource)

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

        # Create 0.orig/include folder
        include_folder = os.path.join(self.case_dir, '0.orig/include')
        if not os.path.exists(include_folder):
            os.makedirs(include_folder)

        start_dir = os.path.join(self.workflows_folder, 'internal/foam')

        # Generate dictionary for rendering p/U files
        field_dict = self._generate_fielddict(att_resource)
        for fieldname in ['p', 'U']:
            source_file = os.path.join(start_dir, 'templates', fieldname)
            dest_file = os.path.join(self.case_dir, '0.orig', fieldname)
            if not self._generate_fieldfile(source_file, dest_file, field_dict):
                self.log().addError('Failed to generate {}'.format(dest_file))
                return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Copy transportProperties file
        source_file = os.path.join(start_dir, 'constant/transportProperties')
        dest_file = os.path.join(self.case_dir, 'constant/transportProperties')
        shutil.copyfile(source_file, dest_file)

        # Generate config files
        self._generate_initialConditions(att_resource)
        self._generate_transportPropertiesConfig(att_resource)

        # Check for 0 directory
        dest_dir = os.path.join(self.case_dir, '0')
        if not os.path.exists(dest_dir):
            source_dir = os.path.join(self.case_dir, '0.orig')
            shutil.copytree(source_dir, dest_dir)

        # 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_allrun_file(
                self.case_dir, solver, copy_casename='snappyHexMesh', restore0=True):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

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

        result = self._create_result()
        return result

    def _copy_system_files(self, att_resource):
        """"""
        start_dir = os.path.join(self.workflows_folder, 'internal/foam')

        # Copy fvSchemes and fvSolution for this application
        app = self.cd_writer.get_application_name(att_resource)
        filenames = ['fvSchemes', 'fvSolution']
        for filename in filenames:
            source_file = os.path.join(start_dir, app, filename)
            if not os.path.exists(source_file):
                self.log().addError('Unable to find {}'.format(source_file))
                continue
            dest_file = os.path.join(self.case_dir, 'system', filename)
            shutil.copyfile(source_file, dest_file)

    def _reset_directory(self):
        """Delete any content from previous jobs."""

        rel_paths = ['0']
        for rel_path in rel_paths:
            path = os.path.join(self.case_dir, rel_path)
            if os.path.exists(path):
                shutil.rmtree(path)

        pathname = '{}/processor*'.format(self.case_dir)
        proc_dirs = glob.glob(pathname)
        for proc_dir in proc_dirs:
            shutil.rmtree(proc_dir)

    def _generate_initialConditions(self, att_resource: smtk.attribute.Resource):
        """Generates initial conditions to include/0.orig directory."""
        kv_list = list()

        pressure_att = att_resource.findAttribute('PressureBoundaryCondition')
        pressure_item = pressure_att.findDouble('Pressure')
        t = ('Pressure', pressure_item.value())
        kv_list.append(t)

        velocity_att = att_resource.findAttribute('VelocityBoundaryCondition')
        velocity_item = velocity_att.findDouble('Velocity')
        if velocity_item.numberOfValues() == 3:  # for backward compatibility
            data = [velocity_item.value(i) for i in range(3)]
        else:
            data = [velocity_item.value(), 0.0, 0.0]
        t = ('FlowVelocity', data)
        kv_list.append(t)

        # Create include directory if needed
        include_dir = os.path.join(self.case_dir, '0.orig/include')
        if not os.path.exists(include_dir):
            os.makedirs(include_dir)

        path = os.path.join(include_dir, 'initialConditions')
        with open(path, 'w') as fp:
            self.cd_writer.write_kvlist(fp, kv_list)

    def _generate_transportPropertiesConfig(self, att_resource: smtk.attribute.Resource):
        """"""
        properties_att = att_resource.findAttribute('PhysicalProperties')
        kviscosity_item = properties_att.findDouble('KinematicViscosity')
        t = ('KinematicViscosity', kviscosity_item.value())
        kv_list = [t]

        # Create include directory if needed
        include_dir = os.path.join(self.case_dir, 'constant/include')
        if not os.path.exists(include_dir):
            os.makedirs(include_dir)

        path = os.path.join(include_dir, 'transportPropertiesConfig')
        with open(path, 'w') as fp:
            self.cd_writer.write_kvlist(fp, kv_list)

    def _generate_fielddict(self, att_resource: smtk.attribute.Resource) -> dict:
        """Uses boundary settings in block mesh to set bottom and front/back data."""
        field_dict = dict()  # return value
        mesh_att = att_resource.findAttribute('BlockMeshBoundaryConditions')

        front_back_item = mesh_att.findString('FrontBackSides')
        front_back_value = front_back_item.value()
        for field in ['p', 'U']:
            key = '{}_front_back'.format(field)
            field_dict[key] = front_back_value

        bottom_item = mesh_att.findString('BottomSide')
        if bottom_item is not None:
            bottom_value = bottom_item.value()
            bottom_type = 'noSlip' if bottom_value == 'wall' else 'zeroGradient'
            field_dict['U_bottom'] = bottom_type

        top_item = mesh_att.findString('TopSide')
        if top_item is not None:
            top_value = top_item.value()
            top_type = 'noSlip' if top_value == 'wall' else 'zeroGradient'
            field_dict['U_top'] = top_type

        return field_dict

    def _generate_fieldfile(self, template_path: str, dest_path: str, field_dict: dict) -> bool:
        """"""
        # Load template file
        template_string = None
        with open(template_path) as fin:
            template_string = fin.read()
        if template_string is None:
            self.log().addError('Failed to find/load template file {}'.format(template_path))
            return False
        dict_template = DictTemplate(template_string)

        # Generate output string and write to dest_path
        output_string = dict_template.safe_substitute(field_dict)
        with open(dest_path, 'w') as fout:
            fout.write(output_string)

        return True

    def _run_process(self, case_dir: str, run_item: smtk.attribute.IntItem, app: str = 'icoFoam') -> bool:
        """"""
        logfilename = '{}.log'.format(app)
        self.logfile = os.path.join(case_dir, logfilename)

        run_mode = run_item.value()
        if run_mode == 'sync':  # run and wait for completion
            args = []
            if hasattr(smtk, 'use_openfoam_docker') and smtk.use_openfoam_docker:
                args = self._docker_run_commands(case_dir)
            args.append(app)

            with open(self.logfile, 'w') as fp:
                completed_proc = subprocess.run(
                    args, stdout=fp, stderr=fp, cwd=case_dir, universal_newlines=True)
                if completed_proc.returncode == 0:
                    self.status = ProcessStatus.Completed
                    success = True

                    # Touch .foam file for visualization
                    foam_filename = '{}.foam'.format(app)
                    foamfile = os.path.join(case_dir, foam_filename)
                    with open(foamfile, 'a'):
                        os.utime(foamfile, 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
            command = '{0} > {1}'.format(app, self.logfile)
            args = []
            if hasattr(smtk, 'use_openfoam_docker') and smtk.use_openfoam_docker:
                args = self._docker_run_commands(case_dir)
            args.append(command)
            proc = subprocess.Popen(args, shell=True, cwd=case_dir, 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
