# ----------------------------------------------------------------------------
# Copyright (c) 2016-2025, 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.
# ----------------------------------------------------------------------------

import collections
import os
import importlib.resources
import re
import shlex
import textwrap
from typing import Any, Callable, Dict, List, Tuple

from qiime2 import ResultCollection
import qiime2.sdk.usage as usage
from qiime2.sdk.usage import (
    UsageVariable, Usage, UsageInputs, UsageOutputs, UsageOutputNames
)
from qiime2.sdk import Action
from qiime2.core.archive.provenance_lib.usage_drivers import (
    build_header, build_footer
)
from qiime2.core.archive.provenance_lib import ProvDAG

import q2cli.util as util
from q2cli.core.state import get_action_state
import q2cli.click.option


def write_example_data(action, output_dir):
    for example_name, example in action.examples.items():
        cli_name = util.to_cli_name(example_name)
        example_path = os.path.join(output_dir, cli_name)

        use = CLIUsage()
        example(use)

        for fn, val in use.get_example_data():
            os.makedirs(example_path, exist_ok=True)
            path = os.path.join(example_path, fn)
            val.save(path)
            try:
                hint = repr(val.type)
            except AttributeError:
                hint = 'Metadata'

            yield hint, path


def write_plugin_example_data(plugin, output_dir):
    for name, action in plugin.actions.items():
        path = os.path.join(output_dir, util.to_cli_name(name))

        yield from write_example_data(action, path)


class CLIUsageVariable(usage.UsageVariable):
    EXT = {
        'artifact': '.qza',
        'artifact_collection': '/',
        'visualization_collection': '/',
        'visualization': '.qzv',
        'metadata': '.tsv',
        'column': '',
        'format': '',
    }

    ELEMENT_EXT = {
        'artifact_collection': EXT['artifact'],
        'visualization_collection': EXT['visualization']
    }

    @property
    def ext(self):
        return self.EXT[self.var_type]

    @staticmethod
    def to_cli_name(val):
        return util.to_cli_name(val)

    def _key_helper(self, input_path, key):
        if self.var_type not in self.COLLECTION_VAR_TYPES:
            raise KeyboardInterrupt(
                f'Cannot key non-collection type {self.var_type}')

        return "%s%s%s" % (input_path, key, self.ELEMENT_EXT[self.var_type])

    def to_interface_name(self):
        if hasattr(self, '_q2cli_ref'):
            return self._q2cli_ref

        cli_name = '%s%s' % (self.name, self.ext)

        # don't disturb file names, this will break importing where QIIME 2
        # relies on specific filenames being present in a dir
        if self.var_type not in ('format', 'column'):
            cli_name = self.to_cli_name(cli_name)

        return cli_name

    def assert_has_line_matching(self, path, expression, key=None):
        if not self.use.enable_assertions:
            return

        INDENT = self.use.INDENT
        input_path = self.to_interface_name()
        expr = shlex.quote(expression)

        if key:
            input_path = self._key_helper(input_path, key)

        lines = [
            'qiime dev assert-result-data %s \\' % (input_path,),
            INDENT + '--zip-data-path %s \\' % (path,),
            INDENT + '--expression %s' % (expr,),
        ]

        self.use.recorder.extend(lines)

    def assert_output_type(self, semantic_type, key=None):
        if not self.use.enable_assertions:
            return

        INDENT = self.use.INDENT
        input_path = self.to_interface_name()

        if key:
            input_path = self._key_helper(input_path, key)

        lines = [
            'qiime dev assert-result-type %s \\' % (input_path,),
            INDENT + '--qiime-type %s' % (str(semantic_type),),
        ]

        self.use.recorder.extend(lines)


class CLIUsage(usage.Usage):
    INDENT = ' ' * 2

    def __init__(self, enable_assertions=False, action_collection_size=None):
        super().__init__()
        self.recorder = []
        self.init_data = []
        self.enable_assertions = enable_assertions
        self.action_collection_size = action_collection_size
        self.output_dir_counter = collections.defaultdict(int)

    def usage_variable(self, name, factory, var_type):
        return CLIUsageVariable(name, factory, var_type, self)

    def render(self, flush=False):
        rendered = '\n'.join(self.recorder)
        if flush:
            self.recorder = []
            self.init_data = []
        return rendered

    def init_artifact(self, name, factory):
        variable = super().init_artifact(name, factory)

        self.init_data.append(variable)

        return variable

    def init_artifact_collection(self, name, factory):
        variable = super().init_artifact_collection(name, factory)

        self.init_data.append(variable)

        return variable

    def construct_artifact_collection(self, name, members):
        variable = super().construct_artifact_collection(
            name, members
        )

        rc_dir = variable.to_interface_name()

        keys = members.keys()
        names = [name.to_interface_name() for name in members.values()]

        keys_arg = '( '
        for key in keys:
            keys_arg += f'{key} '
        keys_arg += ')'
        names_arg = '( '
        for name in names:
            names_arg += f'{name} '
        names_arg += ')'

        lines = [
            '## constructing result collection ##',
            f'rc_name={rc_dir}',
            'ext=.qza',
            f'keys={keys_arg}',
            f'names={names_arg}',
            'construct_result_collection',
            '##',
        ]
        self.recorder.extend(lines)

        return variable

    def get_artifact_collection_member(self, name, variable, key):
        accessed_variable = super().get_artifact_collection_member(
            name, variable, key
        )

        rc_dir = variable.to_interface_name()
        member_fp = os.path.join(rc_dir, f'{key}.qza')

        lines = [
            '## accessing result collection member ##',
            f'ln -s {member_fp} {accessed_variable.to_interface_name()}',
            '##',
        ]
        self.recorder.extend(lines)

        return accessed_variable

    def import_from_format(self, name, semantic_type,
                           variable, view_type=None):
        imported_var = super().import_from_format(
            name, semantic_type, variable, view_type=view_type)

        in_fp = variable.to_interface_name()
        out_fp = imported_var.to_interface_name()

        lines = [
            'qiime tools import \\',
            self.INDENT + '--type %r \\' % (semantic_type,)
        ]

        if view_type is not None:
            if type(view_type) is not str:
                view_type = view_type.__name__
            lines.append(self.INDENT + '--input-format %s \\' % (view_type,))

        lines += [
            self.INDENT + '--input-path %s \\' % (in_fp,),
            self.INDENT + '--output-path %s' % (out_fp,),
        ]

        self.recorder.extend(lines)

        return imported_var

    def init_format(self, name, factory, ext=None):
        if ext is not None:
            name = '%s.%s' % (name, ext.lstrip('.'))

        variable = super().init_format(name, factory, ext=ext)

        self.init_data.append(variable)

        return variable

    def init_metadata(self, name, factory):
        variable = super().init_metadata(name, factory)

        self.init_data.append(variable)

        return variable

    def comment(self, text):
        self.recorder += ['# ' + ln for ln in textwrap.wrap(text, width=74)]

    def peek(self, variable):
        var_name = variable.to_interface_name()
        self.recorder.append('qiime tools peek %s' % var_name)

    def merge_metadata(self, name, *variables):
        var = super().merge_metadata(name, *variables)

        # this is our special "short-circuit" attr to handle special-case
        # .to_interface_name() needs
        var._q2cli_ref = ' '.join(v.to_interface_name() for v in variables)
        return var

    def get_metadata_column(self, name, column_name, variable):
        var = super().get_metadata_column(name, column_name, variable)

        # this is our special "short-circuit" attr to handle special-case
        # .to_interface_name() needs
        var._q2cli_ref = (variable.to_interface_name(), column_name)
        return var

    def view_as_metadata(self, name, variable):
        # use the given name so that namespace behaves as expected,
        # then overwrite it because viewing is a no-op in q2cli
        var = super().view_as_metadata(name, variable)
        # preserve the original interface name of the QZA as this will be
        # implicitly converted to metadata when executed.
        var._q2cli_ref = variable.to_interface_name()
        return var

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

        vars_dict = variables._asdict()

        plugin_name = util.to_cli_name(action.plugin_id)
        action_name = util.to_cli_name(action.action_id)
        self.recorder.append('qiime %s %s \\' % (plugin_name, action_name))

        action_f = action.get_action()
        action_state = get_action_state(action_f)

        ins = inputs.map_variables(lambda v: v.to_interface_name())
        outs = {k: v.to_interface_name() for k, v in vars_dict.items()}
        signature = {s['name']: s for s in action_state['signature']}

        for param_name, value in ins.items():
            self._append_action_line(signature, param_name, value)

        max_collection_size = self.action_collection_size
        if max_collection_size is not None and len(outs) > max_collection_size:
            dir_name = self._build_output_dir_name(plugin_name, action_name)
            self.recorder.append(
                self.INDENT + '--output-dir %s \\' % (dir_name))
            self._rename_outputs(vars_dict, dir_name)
        else:
            for param_name, value in outs.items():
                self._append_action_line(signature, param_name, value)

        self.recorder[-1] = self.recorder[-1][:-2]  # remove trailing \

        return variables

    def _build_output_dir_name(self, plugin_name, action_name):
        base_name = '%s-%s' % (plugin_name, action_name)
        self.output_dir_counter[base_name] += 1
        current_inc = self.output_dir_counter[base_name]
        if current_inc == 1:
            return base_name
        return '%s-%d' % (base_name, current_inc)

    def _rename_outputs(self, vars_dict, dir_name):
        for signature_name, variable in vars_dict.items():
            name = '%s%s' % (signature_name, variable.ext)
            variable._q2cli_ref = os.path.join(dir_name, name)

    def _append_action_line(self, signature, param_name, value):
        param_state = signature[param_name]
        if value is not None:
            for opt, val in self._make_param(value, param_state):
                line = self.INDENT + opt
                if val is not None:
                    line += ' ' + val
                line += ' \\'

                self.recorder.append(line)

    def _make_param(self, value, state):
        state = state.copy()
        type_ = state.pop('type')

        opt = q2cli.click.option.GeneratedOption(prefix=type_[0], **state)
        option = opt.opts[0]

        # INPUTS AND OUTPUTS
        if type_ in ('input', 'output'):
            if isinstance(value, str):
                return [(option, value)]
            else:
                if isinstance(value, set):
                    value = sorted(value)
                return [(option, ' '.join(value))]

        # METADATA FILE
        if state['metadata'] == 'file':
            return [(option, value)]

        # METADATA COLUMN
        if state['metadata'] == 'column':
            # md cols are special, we have pre-computed the interface-specific
            # names and stashed them in an attr, so unpack to get the values
            fn, col_name = value
            return [(option, fn), (opt.q2_extra_opts[0], col_name)]

        # PARAMETERS
        if type_ == 'parameter':
            if isinstance(value, set):
                value = [shlex.quote(str(v)) for v in value]
                return [(option, ' '.join(sorted(value)))]

            if isinstance(value, list):
                return [(option, ' '.join(shlex.quote(str(v)) for v in value))]

            if isinstance(value, (dict, ResultCollection)):
                return [(option, ' '.join(f'{k}:{shlex.quote(str(v))}'
                                          for k, v in value.items()))]

            if type(value) is bool:
                if state['ast']['type'] == 'expression':
                    if value:
                        return [(option, None)]
                    else:
                        return [(opt.secondary_opts[0], None)]
                else:
                    # This is a more complicated param that can't be expressed
                    # as a typical `--p-foo/--p-no-foo` so default to baseline
                    # parameter handling behavior.
                    pass

            if type(value) is str:
                return [(option, shlex.quote(value))]

            return [(option, str(value))]

        raise Exception('Something went terribly wrong!')

    def get_example_data(self):
        for val in self.init_data:
            yield val.to_interface_name(), val.execute()


class ReplayCLIUsageVariable(CLIUsageVariable):
    def to_interface_name(self):
        '''
        Differs from parent method in that metadata is not kebab-cased,
        filepaths are preserved instead.
        '''
        if hasattr(self, '_q2cli_ref'):
            return self._q2cli_ref

        cli_name = '%s%s' % (self.name, self.ext)

        # don't disturb file names, this will break importing where QIIME 2
        # relies on specific filenames being present in a dir
        if self.var_type not in ('format', 'column', 'metadata'):
            cli_name = self.to_cli_name(cli_name)

        return cli_name


class ReplayCLIUsage(CLIUsage):
    shebang = '#!/usr/bin/env bash'
    header_boundary = ('#' * 79)
    set_ex = [
        '# This tells bash to -e exit immediately if a command fails',
        '# and -x show all commands in stdout so you can track progress',
        'set -e -x',
        ''
    ]

    copyright = importlib.resources.read_text(
        f'{__package__}.assets', 'copyright_note.txt'
    ).split('\n')
    how_to = importlib.resources.read_text(
        f'{__package__}.assets', 'cli_howto.txt'
    ).split('\n')

    def __init__(self, enable_assertions=False, action_collection_size=None):
        '''
        Identical to parent but creates header and footer attributes.

        Parameters
        ----------
        enable_assertions : bool
            Whether to render has-line-matching and output type assertions.
        action_collection_size : int
            The number of outputs returned by an action above which outputs are
            grouped into and accessed from an --output-dir.
        '''
        super().__init__()
        self.header = []
        self.footer = []
        self.enable_assertions = enable_assertions
        self.action_collection_size = action_collection_size

    def usage_variable(self, name, factory, var_type):
        return ReplayCLIUsageVariable(name, factory, var_type, self)

    def _append_action_line(self, signature, param_name: str, value):
        '''
        Extends the parent method to accommodate action signatures that may
        differ between those found in provenance and those accessible in the
        currently executing environment.

        Parameters
        ----------
        signature : dict of str -> dict
            Mapping of name of signature item to dict of signature spec data.
        param_name : str
            The name of the parameter for which to render a CLI line.
        value : any
            The value of an item from a Results.
        '''
        param_state = signature.get(param_name)
        if param_state is not None:
            for opt, val in self._make_param(value, param_state):
                line = self.INDENT + opt
                if val is not None:
                    line += ' ' + val
                line += ' \\'
                self.recorder.append(line)
        else:  # no matching param name
            line = self.INDENT + (
                '# FIXME: The following parameter name was not found in '
                'your current\n  # QIIME 2 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')
            cli_name = re.sub('_', '-', param_name)
            line += self.INDENT + '--?-' + cli_name + ' ' + str(value)
            line += ' \\'
            self.recorder.append(line)

    def _make_param(self, value: Any, state: Dict) -> List[Tuple]:
        '''
        Wraps metadata filenames in <> to force users to replace them.

        Parameters
        ----------
        value : any
            A value of a Result.
        state : dict
             A collection of info about an item from an action signature.
             See q2cli.core.state.py.

        Returns
        -------
        list of tuple
            See q2cli.core.usage.CLIUsage._make_param for possible outputs.
        '''
        if state['metadata'] == 'column':
            value = (f'{value[0]}', *value[1:])
        if state['metadata'] == 'file':
            value = f'{value}'
        return super()._make_param(value, state)

    def import_from_format(
        self,
        name: str,
        semantic_type: str,
        variable: UsageVariable,
        view_type: Any = None
    ) -> UsageVariable:
        '''
        Identical to super.import_from_format, but writes
        --input-path <yourdata here> and follows import block with a
        blank line.

        Parameters
        ----------
        name : str
            The name of the created UsageVariable.
        semantic_type : str
            The semantic type of the created UsageVariable.
        variable : UsageVariable
            A UsageVariable of some format type that will materialize the
            data to be imported.
        view_type : format or str
            The view type for importing.

        Returns
        -------
        UsageVariable
            Of type artifact.
        '''
        # need the super().super() here, so pass self to super
        imported_var = Usage.import_from_format(
            self, name, semantic_type, variable, view_type=view_type
        )

        out_fp = imported_var.to_interface_name()

        lines = [
            'qiime tools import \\',
            self.INDENT + '--type %r \\' % (semantic_type,)
        ]

        if view_type is not None:
            lines.append(
                self.INDENT + '--input-format %s \\' % (view_type,)
            )

        lines += [
            self.INDENT + '--input-path <your data here> \\',
            self.INDENT + '--output-path %s' % (out_fp,),
        ]

        lines.append('')
        self.recorder.extend(lines)

        return imported_var

    def init_metadata(
        self, name: str, factory: Callable, dumped_md_fn: str = ''
    ) -> UsageVariable:
        '''
        Like parent, but appropriately handles filepaths for recorded md fps.

        Parameters
        ----------
        name : str
            The name of the UsageVariable to be created.
        factory : Callable
            The factory responsible for generating the realized value of the
            UsageVariable.

        Returns
        -------
        UsageVariable
            Of type metadata.
        '''
        variable = super().init_metadata(name, factory)

        self.init_data.append(variable)

        if dumped_md_fn:
            variable.name = f'"{dumped_md_fn}.tsv"'
        else:
            variable.name = '<your metadata filepath>'

        return variable

    def comment(self, text):
        '''
        Identical to parent method, but pads comments with an extra newline.
        '''
        super().comment(text)
        self.recorder.append('')

    def action(
        self,
        action: Action,
        inputs: UsageInputs,
        outputs: UsageOutputNames
    ) -> UsageOutputs:
        '''
        Overrides parent method to fill in missing outputlines from
        action_f.signature.
        Also pads actions with an extra newline.

        Parameters
        ----------
        action : Action
            The underlying sdk.Action object.
        inputs : UsageInputs
            Mapping of parameter names to arguments for the action.
        outputs : UsageOutputNames
            Mapping of registered output names to usage variable names.

        Returns
        -------
        UsageOutputs
            The results returned by the action.
        '''
        variables = Usage.action(self, action, inputs, outputs)
        vars_dict = variables._asdict()

        # get registered collection of output names so we don't miss any
        action_f = action.get_action()
        missing_outputs = {}
        for output in action_f.signature.outputs:
            try:
                # If we get a match on output-name, the correct pair is already
                # in vars_dict and we can continue
                getattr(variables, output)
            except AttributeError:
                # Otherwise, we should add filler values to missing_outputs
                missing_outputs[output] = f'XX_{output}'

        plugin_name = q2cli.util.to_cli_name(action.plugin_id)
        action_name = q2cli.util.to_cli_name(action.action_id)
        self.recorder.append('qiime %s %s \\' % (plugin_name, action_name))

        action_f = action.get_action()
        action_state = get_action_state(action_f)

        ins = inputs.map_variables(lambda v: v.to_interface_name())
        outs = {k: v.to_interface_name() for k, v in vars_dict.items()}
        outs.update(missing_outputs)
        signature = {s['name']: s for s in action_state['signature']}

        for param_name, value in ins.items():
            self._append_action_line(signature, param_name, value)

        max_collection_size = self.action_collection_size
        if max_collection_size is not None and len(outs) > max_collection_size:
            dir_name = self._build_output_dir_name(plugin_name, action_name)
            self.recorder.append(
                self.INDENT + '--output-dir %s \\' % (dir_name)
            )
            self._rename_outputs(vars_dict, dir_name)
        else:
            for param_name, value in outs.items():
                self._append_action_line(signature, param_name, value)

        self.recorder[-1] = self.recorder[-1][:-2]  # remove trailing `\`

        self.recorder.append('')
        return variables

    def render(self, flush=False):
        '''
        Return a newline-seperated string of CLI script.

        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 CLI code.
        '''
        if self.header:
            self.header = self.header + ['']
        if self.footer:
            self.footer = [''] + self.footer
        rendered = '\n'.join(
            self.header + self.set_ex + self.recorder + self.footer
        )
        if flush:
            self.header = []
            self.footer = []
            self.recorder = []
            self.init_data = []
        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
        ))

        # for creating result collections in bash
        bash_rc_function = [
            'construct_result_collection () {',
            '\tmkdir $rc_name',
            '\ttouch $rc_name.order',
            '\tfor key in "${keys[@]}"; do',
            '\t\techo $key >> $rc_name.order',
            '\tdone',
            '\tfor i in "${!keys[@]}"; do',
            '\t\tln -s ../"${names[i]}" $rc_name"${keys[i]}"$ext',
            '\tdone',
            '}'
        ]
        self.header.extend([
            '## function to create result collections ##',
            *bash_rc_function,
            '##',
        ])

    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))
