# SPDX-FileCopyrightText: 2022-2025 UChicago Argonne, LLC
# SPDX-License-Identifier: MIT

from abc import ABC, abstractmethod
from contextlib import redirect_stdout, redirect_stderr
import os
from pathlib import Path
import shutil
import time
from typing import Optional, List, Union
import uuid

from .database import Database
from .fileutils import cd_tmpdir, PathLike, tee_stdout, tee_stderr, run as run_proc
from .parameters import Parameters
from .results import Results, ExecInfo
from .template import TemplateRenderer
import watts


class Plugin(ABC):
    """Class defining the Plugin interface

    Parameters
    ----------
    extra_inputs
        Extra (non-templated) input files
    show_stdout
        Whether to display output from stdout when :math:`run` is called
    show_stderr
        Whether to display output from stderr when :meth:`run` is called
    unit_system : {'si', 'cgs'}
        Unit system to convert to when rendering input files

    Attributes
    ----------
    plugin_name : str
        Name of the plugin
    unit_system : {'si', 'cgs'}
        Desired system of units for rendering templates

    """

    def __init__(
        self,
        extra_inputs: Optional[List[PathLike]] = None,
        show_stdout: bool = False,
        show_stderr: bool = False,
        unit_system: str = 'si'
    ):
        self.extra_inputs = []
        if extra_inputs is not None:
            self.extra_inputs = [Path(f).resolve() for f in extra_inputs]
        self.show_stdout = show_stdout
        self.show_stderr = show_stderr
        self.unit_system = unit_system

    @abstractmethod
    def prerun(self, params: Parameters):
        ...

    @abstractmethod
    def run(self):
        ...

    @abstractmethod
    def postrun(self, params: Parameters, exec_info: ExecInfo) -> Results:
        ...

    def __call__(
        self,
        params: Parameters = None,
        name: str = '',
        output_dir: Optional[PathLike] = None,
        verbose: bool = True,
        cleanup: bool = True,
        **kwargs
    ) -> Results:
        """Run the complete workflow for the plugin

        Parameters
        ----------
        params
            Parameters used in generating inputs
        name
            Name associated with execution of plugin
        output_dir
            Relative path of the directory created in the database. If None, a
            unique directory name is created using the uuid module.
        verbose
            Whether to print execution information
        cleanup
            Determines whether the temporary directory will be cleaned up
            immediately after execution.
        **kwargs
            Keyword arguments passed to the `run` method

        Returns
        -------
        Results from running plugin
        """
        db = Database()
        plugin_name = self.plugin_name
        if verbose:
            if name:
                print(f'[watts] Calling {plugin_name} ({name})...')
            else:
                print(f'[watts] Calling {plugin_name}...')

        # Generate empty Parameters object if none provided
        if params is None:
            params = Parameters()

        # Create execution info
        timestamp = time.time_ns()
        exec_info = ExecInfo(db.job_id, plugin_name, name, timestamp)

        with cd_tmpdir(cleanup=cleanup):
            # Copy extra inputs to execution directory
            cwd = Path.cwd()
            for path in self.extra_inputs:
                shutil.copy(str(path), str(cwd))  # Remove str() for Python 3.8+

            # Generate input files and perform any other prerun actions
            self.prerun(params)

            # Execute the code, redirecting stdout/stderr if requested
            with open(f'{plugin_name}_log.txt', 'w') as outfile:
                func_stdout = tee_stdout if self.show_stdout else redirect_stdout
                func_stderr = tee_stderr if self.show_stderr else redirect_stderr
                with func_stdout(outfile), func_stderr(outfile):
                    self.run(**kwargs)

            # Collect results and perform any postrun actions
            result = self.postrun(params, exec_info)

            # Create new directory for results and move files there
            if output_dir is None:
                output_dir = uuid.uuid4().hex
            workflow_path = db.path / output_dir
            workflow_path.mkdir()
            try:
                result.move_files(workflow_path)
            except Exception:
                # If error occurred, make sure we remove results directory so it
                # doesn't pollute database
                #shutil.rmtree(workflow_path)
                raise

        # Add result to database
        db.add_result(result)

        return result


class PluginGeneric(Plugin):
    """Plugin that relies on generating a template file

    This class can be used to control the execution of an arbitrary executable,
    first rendering one or more templated input files.

    Parameters
    ----------
    executable
        Path to executable
    execute_command
        List of command-line arguments, where each is formatted using the
        instance of the class as ``self``. The first string normally indicates
        the executable, i.e. "{self.executable}". The rendered input file can be
        accessed as "{self.input_name}". A single string of command-line
        arguments is also accepted.
    template_file
        Path to template file
    extra_inputs
        Extra (non-templated) input files
    extra_template_inputs
        Extra templated input files
    plugin_name
        Name of the plugin
    show_stdout
        Whether to display output from stdout when :math:`run` is called
    show_stderr
        Whether to display output from stderr when :meth:`run` is called
    unit_system : {'si', 'cgs'}
        Unit system to convert to when rendering input files

    Attributes
    ----------
    executable
        Path to plugin executable
    execute_command
        List of command-line arguments used to call the executable

    """
    def __init__(self,
        executable: PathLike,
        execute_command: Union[List[str], str],
        template_file: PathLike,
        extra_inputs: Optional[List[PathLike]] = None,
        extra_template_inputs: Optional[List[PathLike]] = None,
        plugin_name: str = 'Generic',
        show_stdout: bool = False,
        show_stderr: bool = False,
        unit_system: str = 'si',
    ):
        super().__init__(extra_inputs, show_stdout, show_stderr, unit_system)
        self.plugin_name = plugin_name
        self.render_template = TemplateRenderer(template_file)
        self.extra_render_templates = []
        self.input_name = 'input_rendered'
        if extra_template_inputs is not None:
            self.extra_render_templates = [TemplateRenderer(f, '') for f in extra_template_inputs]

        self.executable = executable
        if isinstance(execute_command, str):
            self._execute_command = execute_command.split()
        else:
            self._execute_command = execute_command

    @property
    def executable(self) -> Path:
        return self._executable

    @executable.setter
    def executable(self, exe: PathLike):
        if shutil.which(exe) is None:
            raise RuntimeError(f"{self.plugin_name} executable '{exe}' is missing.")
        self._executable = Path(exe)

    @property
    def execute_command(self) -> List[str]:
        return [item.format(self=self) for item in self._execute_command]

    def prerun(self, params: Parameters, filename: Optional[str] = None):
        """Render the template based on model parameters

        Parameters
        ----------
        params
            Parameters used to render template
        filename
            Filename for rendered template
        """
        # If the 'input_name' attribute is set, use that as default when
        # filename is not explicitly passed
        if filename is None and self.input_name is not None:
            filename = self.input_name

        # Make a copy of params and convert units if necessary -- the original
        # params remains unchanged
        params_copy = params.convert_units(system=self.unit_system)

        # Render the template
        self.render_template(params_copy, filename=filename)
        for render_template in self.extra_render_templates:
            render_template(params_copy)

    def postrun(self, params: Parameters, exec_info: ExecInfo, **kwargs) -> Results:
        """Read simulation results and create results object

        Parameters
        ----------
        params
            Parameters used to generate input files
        exec_info
            Execution information
        **kwargs
            Keyword arguments for Results subclasses

        Returns
        -------
        Results object
        """

        # Determine inputs and outputs
        inputs = [self.input_name] + [p.name for p in self.extra_inputs]
        for renderer in self.extra_render_templates:
            inputs.append(renderer.template_file.name)
        outputs = [p for p in Path.cwd().iterdir() if p.name not in inputs]

        # Get correct Results subclass and return instance
        results_cls = getattr(watts, f'Results{self.plugin_name}', Results)
        return results_cls(params, exec_info, inputs, outputs, **kwargs)

    def run(self, mpi_args: Optional[List[str]] = None,
            extra_args: Optional[List[str]] = None):
        """Run plugin

        Parameters
        ----------
        mpi_args
            MPI execute command and any additional MPI arguments to pass,
            e.g. ['mpiexec', '-n', '8'].
        extra_args
            Additional command-line arguments to append after the main command

        """
        if mpi_args is None:
            mpi_args = []
        if extra_args is None:
            extra_args = []
        run_proc(mpi_args + self.execute_command + extra_args)


def _find_executable(path: PathLike, environment_variable: str) -> Path:
    """Determine executable for a given code with a hint from environment variable

    Parameters
    ----------
    path
        Name of exeuctable or absolute path
    environment_variable
        Environment variable indicating directory where executable is found

    Returns
    -------
    Path to executable
    """

    exe = Path(path)

    # If path is already an absolute path, use it
    if exe.is_absolute():
        return exe

    # Check for environment variable
    base_dir = os.environ.get(environment_variable)
    return base_dir / exe if base_dir is not None else exe
