from typing import Dict, List, Tuple, Union

import traitlets
import yaml
from nbconvert.preprocessors import Preprocessor
from nbformat import NotebookNode

from ._nested_dict_updater import nested_dict_updater
from ._patterns import (
    build_prefixed_regex_pattern,
    build_prefixed_regex_pattern_with_value,
)

__all__ = [
    "MetaDataListInjectorPreprocessor",
    "MetaDataMapInjectorPreprocessor",
]


class MetaDataListInjectorPreprocessor(Preprocessor):
    """
    Parse all *code* cells and append the matched magic comments with the
    `prefix` to the `metadata_group` list.
    These `strings` must be on their own line and only contain the `prefix`,
    a `string` from `strings` (i.e., the magic comment) and whitespace characters.
    """

    metadata_group: str = traitlets.Unicode(default_value="tags").tag(config=True)
    """Metadata group to which the matched magic comment will be appended to if
    it doesn't already exist. Default is `tags`."""
    strings: List[str] = traitlets.List(traitlets.Unicode(), minlen=1).tag(config=True)
    """List of strings (magic comments) that define the text that will be matched and
    injected into the selected metadata group."""
    prefix: str = traitlets.Unicode(default_value="#").tag(config=True)
    """The prefix that indicates the possible start of a magic comment line.
    Should be comment character of the language. By default `#`."""
    remove_line: bool = traitlets.Bool(default_value=True).tag(config=True)
    """By default remove the matching line in the code-cell."""

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        if self.metadata_group == "":
            raise ValueError("metadata_group myst be non-empty string!")

    def _write_tag(self, tag: str, cell: NotebookNode) -> NotebookNode:
        tags = cell.setdefault("metadata", {}).setdefault(self.metadata_group, [])
        if not isinstance(tags, list):
            raise RuntimeError(
                f"Trying to to set metadata list but entry {self.metadata_group} is of type: {type(tags)}",
                self.metadata_group,
            )
        if tag not in tags:
            tags.append(tag)
        cell["metadata"][self.metadata_group] = tags
        return cell

    def preprocess_cell(
        self, cell: NotebookNode, resource: Union[Dict, None], _index: int
    ) -> Tuple[NotebookNode, Union[Dict, None]]:
        """Inject metadata to code-cell if match is found"""
        if cell["cell_type"] == "markdown":
            return cell, resource
        for string in self.strings:
            pattern = build_prefixed_regex_pattern(prefix=self.prefix, key_term=string)
            matches = pattern.finditer(cell.source)
            for m in matches:
                tag = m.group("key")
                cell = self._write_tag(tag, cell)
            if self.remove_line:
                cell.source = pattern.sub("", cell.source)
        return cell, resource


class MetaDataMapInjectorPreprocessor(Preprocessor):
    """
    Parse all *code* cells and add the matched key-value pairs with the
    `prefix` to the `metadata_group` dictionary.
    The key-value pairs are generated by searching for each `key` of `keys` followed
    by `delimiter` and the value.
    """

    metadata_group: str = traitlets.Unicode().tag(config=True)
    """Metadata group into which the matched key-value pairs will be written."""
    keys: List[str] = traitlets.List(traitlets.Unicode()).tag(config=True)
    """List of keys that will be used as a key for the `metadata_group` dictionary entry and is followed by the `delimiter` and `value`."""
    prefix: str = traitlets.Unicode(default_value="#").tag(config=True)
    """The prefix that indicates the possible start of a magic comment line. Should be comment character of the language."""
    remove_line: bool = traitlets.Bool(default_value=True).tag(config=True)
    """By default remove the matching line in the code-cell."""
    delimiter: str = traitlets.Unicode(default_value="=").tag(config=True)
    """Delimiter that separates the key from the value."""
    value_to_yaml: bool = traitlets.Bool(default_value=False).tag(config=True)
    """Parse the value as yaml syntax before writing it as a dictionary. Default is `False`."""
    allow_nested_keys: bool = traitlets.Bool(default_value=False).tag(config=True)
    """Allow defining nested key access for nested setting of metadata group. Access next level with `.` (hardcoded)"""

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        if self.metadata_group == "":
            raise ValueError("metadata_group myst be non-empty string!")

    def preprocess_cell(
        self, cell: NotebookNode, resource: Union[Dict, None], _index: int
    ) -> Tuple[NotebookNode, Union[Dict, None]]:
        """Inject metadata dict entry to code-cell if match is found"""
        if cell["cell_type"] == "markdown":
            return cell, resource
        for key_term in self.keys:
            pattern = build_prefixed_regex_pattern_with_value(
                prefix=self.prefix,
                key_term=key_term,
                delimiter=self.delimiter,
                expand_key_term=self.allow_nested_keys,
            )
            matches = pattern.finditer(cell.source)
            for m in matches:
                value = m.group("value")
                parsed_value = yaml.safe_load(value) if self.value_to_yaml else value
                key = m.group("key")  # Could be expanded key!
                keys = key.split(".") if self.allow_nested_keys else [key]
                entry = cell.setdefault("metadata", {}).setdefault(
                    self.metadata_group, {}
                )
                nested_dict_updater(entry, keys, parsed_value)
            if self.remove_line:
                cell.source = pattern.sub("", cell.source)
        return cell, resource


# class GlobalMetaDataInjectorPreprocessor(Preprocessor):
#     """
#     Parse all *code* cells and convert the matching `prefix` `key` `value`
#     lines to the global `metadata` field.

#     To clean up the output, the lines containing any `string` may be removed
#     by setting `remove_line=True` (default).

#     The provided list of `keys` will be used to access the *global* `metadata` field
#     and insert the value that is followed by the `key` in the code cell.
#     Note that the global metadata field will be overwritten if multiple cells define the
#     field's value.
#     """

#     keys = List(Unicode()).tag(config=True)
#     prefix = Unicode(default_value="#").tag(config=True)
#     delimiter = Unicode(default_value=r"=").tag(config=True)

#     def preprocess(self, nb, resources):
#         if len(self.keys) == 0:
#             return nb, resources

#         for cell in nb.cells:
#             if cell["cell_type"] == "markdown":
#                 continue
#             for key in self.keys:
#                 pattern = build_prefixed_regex_pattern_with_value(
#                     self.prefix, key, delimiter=self.delimiter
#                 )
#                 m = pattern.search(cell.source)
#                 if m is not None:
#                     value = m.group("value")
#                     nb.setdefault("metadata", {})
#                     nb["metadata"][key] = value
#         return nb, resources
