# ----------------------------------------------------------------------------
# Copyright (c) 2016-2026, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------
from datetime import datetime
from importlib.metadata import metadata
import importlib.resources
import textwrap
from typing import Any, Callable, List, Type

from .parse import ProvDAG, ProvNode

from rachis.sdk import Action
from rachis.plugins import ArtifactAPIUsage
from rachis.sdk.usage import (
    Usage, UsageAction, UsageVariable, UsageInputs, UsageOutputs,
)


def build_header(
    shebang: str = '',
    boundary: str = '',
    copyright: str = '',
    extra_text: List[str] = []
) -> List[str]:
    '''
    Constructs the header contents for a replay script.

    Parameters
    ----------
    shebang : str
        The shebang line to add to the rendered script, if any.
    boundary : str
        The visual boundary to add to the rendred script, if any.
    copyright : str
        The copyright notice to add to the rendered script, if any.
    extra_text : list of str
        Extra lines of text to add to the header in the rendered script,
        if any.

    Returns
    -------
    list of str
        The constructed header lines.
    '''
    rachis_md = metadata('rachis')
    vzn = rachis_md['Version']
    ts = datetime.now()
    header = []

    if shebang:
        header.append(shebang)

    if boundary:
        header.append(boundary)

    header.extend([
        f'# Auto-generated by Rachis v.{vzn} at '
        f'{ts.strftime("%I:%M:%S %p")} on {ts.strftime("%d %b, %Y")}',
    ])

    if copyright:
        header.extend(copyright)

    header.append(
        '# For User Support, post to the QIIME 2 Forum at '
        'https://forum.qiime2.org.'
    )

    if extra_text:
        header.extend(extra_text)
    if boundary:
        header.append(boundary)
    return header


def build_footer(dag: ProvDAG, boundary: str) -> List[str]:
    '''
    Constructs the footer contents for a replay script.

    Parameters
    ----------
    dag : ProvDAG
        The ProvDAG object representing the input artifact(s).
    boundary : str
        The visual boundary demarcating the footer.

    Returns
    -------
    list of str
        The constructed footer lines as a list of strings.
    '''
    footer = []
    pairs = []
    uuids = sorted(dag._parsed_artifact_uuids)
    # two UUIDs fit on a line
    for idx in range(0, u_len := len(uuids), 2):
        if idx == u_len - 1:
            pairs.append(f'# {uuids[idx]}')
        else:
            pairs.append(f'# {uuids[idx]} \t {uuids[idx + 1]}')

    footer.append(boundary)
    footer.append(
        '# The following Rachis Results were parsed to produce this script:'
    )
    footer.extend(pairs)
    footer.append(boundary)
    footer.append('')
    return footer


class ReplayUsageAction(UsageAction):
    def __init__(self, plugin_id: str, action_id: str, node: ProvNode):
        self.node = node
        if node._action_present_:
            super().__init__(plugin_id, action_id)
        else:
            if plugin_id == '':
                raise ValueError('Must specify a value for plugin_id.')

            if action_id == '':
                raise ValueError('Must specify a value for action_id.')

            self.plugin_id: str = plugin_id
            self.action_id: str = action_id

    def get_action(self) -> Action | None:
        if self.node._action_present_:
            return super().get_action()

        return None


class ReplayPythonUsage(ArtifactAPIUsage):
    UsageAction: Type[ReplayUsageAction] = ReplayUsageAction

    shebang = '#!/usr/bin/env python'
    header_boundary = '# ' + ('-' * 77)
    copyright = importlib.resources.read_text(
        'rachis.core.archive.provenance_lib.assets',
        'copyright_note.txt'
        ).split('\n')
    how_to = importlib.resources.read_text(
        'rachis.core.archive.provenance_lib.assets',
        'python_howto.txt'
        ).split('\n')

    def __init__(
        self,
        enable_assertions: bool = False,
        action_collection_size: int = 2
    ):
        '''
        Identical to parent, but with smaller default action_collection_size.

        Parameters
        ----------
        enable_assertions : bool
            Whether to render has-line-matching and output type assertions.
        action_collection_size : int
            The maximum number of outputs returned by an action above which
            results are grouped into and destructured from a single variable.
        '''
        super().__init__()
        self.enable_assertions = enable_assertions
        self.action_collection_size = action_collection_size
        self._reset_state(reset_global_imports=True)

    def _reset_state(self, reset_global_imports=False):
        '''
        Clears all state associated with the usage driver, excepting global
        imports by default.

        Parameters
        ----------
        resest_global_imports : bool
            Whether to reset self.global_imports to an empty set.
        '''
        self.local_imports = set()
        self.header = []
        self.recorder = []
        self.footer = []
        self.init_data_refs = dict()
        if reset_global_imports:
            self.global_imports = set()

    def _template_action(
        self, action: Action, input_opts: UsageInputs, variables: UsageOutputs
    ):
        '''
        Templates the artifact api python code for the action `action`.

        Extends the parent method to:
        - accommodate action signatures that may differ between those found in
          provenance and those accessible in the currently executing
          environment
        - render artifact api code that saves the results to disk.

        Parameters
        ----------
        action : Action
            The Rachis Action object.
        input_opts : UsageInputs
            The UsageInputs mapping for the action.
        variables : UsageOutputs
            The UsageOutputs object for the action.
        '''
        action_f = action.get_action()
        if (
            len(variables) > self.action_collection_size
            or len(action_f.signature.outputs) > 5
        ):
            output_vars = 'action_results'
        else:
            output_vars = self._template_outputs(action, variables)

        plugin_id = action.plugin_id
        action_id = action.action_id
        lines = [
            f'{output_vars} = {plugin_id}_actions.{action_id}('
        ]

        all_inputs = (list(action_f.signature.inputs.keys()) +
                      list(action_f.signature.parameters.keys()))
        for k, v in input_opts.items():
            line = ''
            if k not in all_inputs:
                line = self.INDENT + (
                    '# FIXME: The following parameter name was not found in '
                    'your current\n    # Rachis environment. This may occur '
                    'when the plugin version you have\n    # installed does '
                    'not match the version used in the original analysis.\n '
                    ' # Please see the docs and correct the parameter name '
                    'before running.\n'
                )
            line += self._template_input(k, v)
            lines.append(line)

        lines.append(')')

        if (
            len(variables) > self.action_collection_size
            or len(action.get_action().signature.outputs) > 5
        ):
            for k, v in variables._asdict().items():
                interface_name = v.to_interface_name()
                lines.append('%s = action_results.%s' % (interface_name, k))

        lines.append(
            '# SAVE: comment out the following with \'# \' to skip saving '
            'Results to disk'
        )
        for k, v in variables._asdict().items():
            interface_name = v.to_interface_name()
            lines.append(
                '%s.save(\'%s\')' % (interface_name, interface_name,))

        lines.append('')
        self._add(lines)

    def _template_action_not_found(
        self, action: Action, input_opts: UsageInputs, variables: UsageOutputs
    ):
        '''
        Templates the artifact api python code for the action `action`.

        Extends the parent method to:
        - accommodate action signatures that may differ between those found in
          provenance and those accessible in the currently executing
          environment
        - render artifact api code that saves the results to disk.

        Parameters
        ----------
        action : Action
            The Rachis Action object.
        input_opts : UsageInputs
            The UsageInputs mapping for the action.
        variables : UsageOutputs
            The UsageOutputs object for the action.
        '''
        output_vars = 'action_results'

        plugin_id = action.plugin_id
        action_id = action.action_id
        lines = [
            '# FIXME: The following action was not found in your current '
            'QIIME 2\n# environment. Please ensure the action and its '
            'parameters are correct before\n# running.',
            f'{output_vars} = {plugin_id}_actions.{action_id}('
        ]

        for k, v in input_opts.items():
            line = self._template_input(k, v)
            lines.append(line)

        lines.append(')')
        for k, v in variables._asdict().items():
            interface_name = v.to_interface_name()
            lines.append('%s = action_results.%s' % (interface_name, k))

        lines.append(
            '# SAVE: comment out the following with \'# \' to skip saving '
            'Results to disk'
        )
        for k, v in variables._asdict().items():
            interface_name = v.to_interface_name()
            lines.append(
                '%s.save(\'%s\')' % (interface_name, interface_name,))

        lines.append('')
        self._add(lines)

    def _template_outputs(
        self, action: Action, variables: UsageOutputs
    ) -> str:
        '''
        Extends the parent method to allow the replay an action when a
        ProvDAG doesn't have a record of all outputs from an action. These
        unknown outputs are given the conventional '_' variable name.

        Parameters
        ----------
        action : Action
            The Action object associated with the output variables.
        variables : UsageOutputs
            The UsageOutputs object associated with the action.

        Returns
        -------
        str
            The templated output variables names as a comma-separated string.
        '''
        output_vars = []
        action_f = action.get_action()

        # need to coax the outputs into the correct order for unpacking
        for output in action_f.signature.outputs:
            try:
                variable = getattr(variables, output)
                output_vars.append(str(variable.to_interface_name()))
            except AttributeError:
                output_vars.append('_')

        if len(output_vars) == 1:
            output_vars.append('')

        return ', '.join(output_vars).strip()

    def action(self, action, inputs, outputs):
        if action.node._action_present_:
            variables = super().action(action, inputs, outputs)
        else:
            variables = self.action_not_found(action, inputs, outputs)

        return variables

    def action_not_found(self, action, inputs, outputs):
        variables = super().action_not_found(action, inputs, outputs)

        self._plugin_import_as_name_not_found(action)

        inputs = inputs.map_variables(lambda v: v.to_interface_name())
        self._template_action_not_found(action, inputs, variables)

        return variables

    def _plugin_import_as_name_not_found(self, action):
        base = f'rachis.plugins.{action.plugin_id}.actions'
        as_ = f'{action.plugin_id}_actions'
        self._update_imports(
            import_='%s.actions' % (base,), as_=as_, missing_=True
        )
        return as_

    def init_metadata(
        self, name: str, factory: Callable, dumped_md_fn: str = ''
    ) -> UsageVariable:
        '''
        Renders the loading of Metadata from disk.

        Parameters
        ----------
        name : str
            The name of the created and returned UsageVariable.
        factory : Callable
            The factory responsible for constructing the metadata
            UsageVariable.
        dumped_md_fn : str
            Optional. The filename of the dumped metadata.

        Returns
        -------
        UsageVariable
            The UsageVariable of var_type metadata corresponding to the
            loaded metadata.
        '''
        var = super().init_metadata(name, factory)
        self._update_imports(from_='rachis', import_='Metadata')
        input_fp = var.to_interface_name()
        if dumped_md_fn:
            lines = [f'{input_fp} = Metadata.load("{dumped_md_fn}.tsv")']
        else:
            self.comment(
                'NOTE: You may substitute already-loaded Metadata for the '
                'following, or cast a pandas.DataFrame to Metadata as needed.'
            )
            lines = [f'{input_fp} = Metadata.load(<your metadata filepath>)']

        self._add(lines)
        return var

    def import_from_format(
        self,
        name: str,
        semantic_type: str,
        variable: UsageVariable,
        view_type: Any = None
    ):
        '''
        Extends the parent method to:
        - use '<your data here>' instead of import_fp for the import filepath
        - render artifact api code that saves the result to disk.

        Parameters
        ----------
        name : str
            The name of the UsageVariable `variable`.
        semantic_type : str
            The semantic type of the UsageVariable `variable`.
        variable : UsageVariable
            The usage variable object.
        view_type : str or some format
            The view type to use for importing.

        Returns
        -------
        UsageVariable
            The imported artifact UsageVariable.
        '''
        imported_var = Usage.import_from_format(
            self, name, semantic_type, variable, view_type=view_type
        )

        interface_name = imported_var.to_interface_name()
        import_fp = self.repr_raw_variable_name('<your data here>')

        lines = [
            '%s = Artifact.import_data(' % (interface_name,),
            self.INDENT + '%r,' % (semantic_type,),
            self.INDENT + '%r,' % (import_fp,),
        ]

        if view_type is not None:
            if type(view_type) is not str:
                # Show users where these formats come from when used in the
                # Python API to make things less 'magical'.
                import_path = super()._canonical_module(view_type)
                view_type = view_type.__name__
                if import_path is not None:
                    self._update_imports(from_=import_path,
                                         import_=view_type)
                else:
                    # May be in scope already, but something is quite wrong at
                    # this point, so assume the plugin_manager is sufficiently
                    # informed.
                    view_type = repr(view_type)
            else:
                view_type = repr(view_type)

            lines.append(self.INDENT + '%s,' % (view_type,))

        lines.extend([
            ')',
            '# SAVE: comment out the following with \'# \' to skip saving this'
            ' Result to disk',
            '%s.save(\'%s\')' % (interface_name, interface_name,),
            ''
        ])

        self._update_imports(from_='rachis', import_='Artifact')
        self._add(lines)

        return imported_var

    class repr_raw_variable_name:
        # allows us to repr col name without enclosing quotes
        # (as in rachis.plugins.ArtifactAPIUsageVariable)
        def __init__(self, value):
            self.value = value

        def __repr__(self):
            return self.value

    def comment(self, line: str):
        '''
        Communicate that a comment should be rendered.

        Parameters
        ----------
        line : str
            The comment to be rendered.
        '''
        LINE_LEN = 79
        lines = textwrap.wrap(
            line,
            LINE_LEN,
            break_long_words=False,
            initial_indent='# ',
            subsequent_indent='# '
        )
        lines.append('')
        self._add(lines)

    def render(self, flush: bool = False) -> str:
        '''
        Return a newline-seperated string of Artifact API python code.

        Parameters
        ----------
        flush : bool
            Whether to 'flush' the current code. Importantly, this will clear
            the top-line imports for future invocations.

        Returns
        -------
        str
            The rendered string of python code.
        '''
        sorted_imps = sorted(self.local_imports)
        if self.header:
            self.header = self.header + ['']
        if self.footer:
            self.footer = [''] + self.footer
        if sorted_imps:
            sorted_imps = sorted_imps + ['']
        rendered = '\n'.join(
            self.header + sorted_imps + self.recorder + self.footer
        )
        if flush:
            self._reset_state()
        return rendered

    def build_header(self):
        '''Constructs a renderable header from its components.'''
        self.header.extend(build_header(
            self.shebang, self.header_boundary, self.copyright, self.how_to
        ))

    def build_footer(self, dag: ProvDAG):
        '''
        Constructs a renderable footer using the terminal uuids of a ProvDAG.
        '''
        self.footer.extend(build_footer(dag, self.header_boundary))
