# Copyright (c) 2022 Chris Duf <chris@openmarl.org>
#
# SPDX-License-Identifier: Apache-2.0

"""Built-in 'find' command."""


import re

from abc import abstractmethod
from typing import Tuple, List, Union

from devicetree.edtlib import Node

from dtsh.dtsh import DtshCommand, DtshCommandOption, Dtsh, DtshAutocomp, DtshVt
from dtsh.dtsh import DtshCommandFlagLongFmt, DtshCommandArgLongFmt
from dtsh.dtsh import DtshCommandUsageError
from dtsh.tui import LsNodeTable, DtshTui


class FindCriterion(object):
    """Interface for search criteria.
    """

    # Plain text pattern (used only for plain text search, w/o wildcard)
    _pattern: str

    def __init__(self, pattern: str) -> None:
        """Initialize criterion.

        Arguments:
        pattern - the search pattern
        """
        self._pattern = pattern

    @abstractmethod
    def match(self, node: Node) -> bool:
        """Returns True if the node does match this criterion.
        """


class DtshBuiltinFind(DtshCommand):
    """Find devicetree nodes.

DESCRIPTION
The `find` command will search `PATH` for devicetree nodes, where `PATH` is either:

- an absolute path to a devicetree node, e.g. `/soc`
- a relative path to a devicetree node, e.g. `soc`

`PATH` supports simple path substitution:

- a leading `.` is interpreted as the current working node
- a leading `..` is interpreted as the current working node's parent

If `PATH` is unspecified, `find` will search the current working node.

The `find` command always search recursively.

The search supports multiple criteria:

- **--name**: match nodes by names
- **--compat**: match nodes by compatible strings
- **--bus**: match nodes by *bus device*
- **--interrupt**: match nodes by interrupts
- **--enabled-only**: search only enabled (aka *okay*) nodes

Criteria are additive (logical `AND`), such that:

    find --bus * --interrupt *

will match all bus devices that generate IRQs.

Without any criterion, the `find` command will behave like its
POSIX homonym: the search will match all nodes at `PATH`.

Criteria are defined in the next sections.

By default, `find` will only print the found node paths: use the **-l** option to
enable a more detailed (aka *rich*) output.

The **-f <fmt>** option allows to specify the visible columns with a
format string.

Valid column specifiers are:

        | Specifier | Format                                    | DTSpec  |
        |-----------|-------------------------------------------|---------|
        | `N`       | The node name                             | 2.2.1   |
        | `a`       | The unit-address                          |         |
        | `n`       | The node name with the address striped    |         |
        | `d`       | The description from the node binding     |         |
        | `p`       | The node path name                        | 2.2.3   |
        | `l`       | The node 'label' property                 |         |
        | `L`       | All known labels for the node             |         |
        | `s`       | The node 'status' property                | 2.3.4   |
        | `c`       | The 'compatible' property for the node    | 2.3.1   |
        | `C`       | The node binding (aka matched compatible) |         |
        | `A`       | The node aliases                          |         |
        | `b`       | The bus device information for the node   |         |
        | `r`       | The node 'reg' property                   | 2.3.6   |
        | `i`       | The interrupts generated by the node      | 2.4.1.1 |

Use the **-c** options to count the matched nodes.
Add the *quiet* flag (**-q**) to only print the nodes count,
not the nodes themselves.
 
Set the **--pager** option to page the command's output using the system pager.

**Search by name**

Match nodes by the `node-name` component
([DTSpec 2.2.1](https://devicetree-specification.readthedocs.io/en/stable/devicetree-basics.html#node-names)).

If the search pattern does not contain any wildcard,
defaults to plain text search.

If the search pattern contains `*` wildcards,
a regular expression type search is assumed
where the wildcards are replaced by a character class appropriate for node names.

Then:

- `*` will match any node-name
- `<prefix>*` will match a node-name if it starts with `<prefix>`
- `*<postfix>` will match a node-name if it ends with `<postfix>`
- `<prefix>*<postfix>` will match a node-name if it starts with `<prefix>`
  and ends with `<postfix>`

**Search by compatible**

Match nodes by `compatible` string
([DTSpec 2.3.1](https://devicetree-specification.readthedocs.io/en/stable/devicetree-basics.html#compatible)).

If the search pattern does not contain any wildcard,
defaults to plain text search.

If the search pattern contains `*` wildcards,
a regular expression type search is assumed
where the wildcards are replaced by a character class appropriate for
compatible strings.

Then:

- `*` will match any compatible string
- `<prefix>*` will match a compatible string if it starts with `<prefix>`
- `*<postfix>` will match a compatible string if it ends with `<postfix>`
- `<prefix>*<postfix>` will match a compatible string if it starts with
   `<prefix>` and ends with `<postfix>`

**Search by bus device**

Match nodes that provide or appear on a *bus device*, e.g. `i2c` or `spi`.

If the search pattern does not contain any wildcard,
defaults to plain text search.

If the search pattern contains `*` wildcards,
a regular expression type search is assumed
where the wildcards are replaced by a character class appropriate for
bus devices.

Then:

- `*` will match any bus device
- `<prefix>*` will match a bus device if it starts with `<prefix>`
- `*<postfix>` will match a bus device if it ends with `<postfix>`
- `<prefix>*<postfix>` will match a bus device if it starts with `<prefix>`
  and ends with `<postfix>`

**Search by IRQ**

Match nodes by interrupt numbers or names.

If the search pattern successfully converts to an integer,
a search by IRQ number is assumed.

If the conversion fails, a search by interrupt name is assumed.

If the pattern equals to `*`, the search will match all interrupts: this permits
to *find* all nodes that generate IRQs.

If the search pattern does not contain any wildcard,
defaults to plain text search.

If the search pattern contains `*` wildcards,
a regular expression type search is assumed
where the wildcards are replaced by a character class appropriate for IRQ names.

Then:

- `*` will match any IRQ name
- `<prefix>*` will match an IRQ name if it starts with `<prefix>`
- `*<postfix>` will match an IRQ name if it ends with `<postfix>`
- `<prefix>*<postfix>` will match an IRQ name if it starts with `<prefix>`
  and ends with `<postfix>`

EXAMPLES

1. All nodes

Without any criterion, the search will match all devicetree nodes.

Count the devicetree nodes:

```
/
❯ find -cq

132 nodes.
```

Dump the enabled nodes:

```
/
❯ find -c --enabled-only
/
/chosen
/aliases
/soc
/soc/interrupt-controller@e000e100
/soc/ficr@10000000
/soc/uicr@10001000
...
/buttons/button_3
/connector
/analog-connector

119 nodes.
```

2. Find nodes by name

To find nodes which name contains `gpio` (*plain text* search):

```
/
❯ find --name gpio -l
Path                  Aliases  Labels
─────────────────────────────────────
/soc/gpiote@40006000           gpiote
/soc/gpio@50000000             gpio0
/soc/gpio@50000300             gpio1
```

To find nodes which name ends with `gpio` (RE search):

```
/
❯ find --name *gpio -l
Path                  Aliases  Labels
─────────────────────────────────────
/soc/gpio@50000000             gpio0
/soc/gpio@50000300             gpio1
```

3. Find nodes by compatible strings

To find nodes that may involve a TWI *compatible* driver:

```
❯ find --compat twi  -l
Path               Compatible      Description
──────────────────────────────────────────────────────────────────────
/soc/i2c@40003000  nordic,nrf-twi  Nordic nRF family TWI (TWI master)…
/soc/i2c@40004000  nordic,nrf-twi  Nordic nRF family TWI (TWI master)…
```

4. Find nodes by bus devices

To find all enabled bus devices:

```
/
❯ find --enabled-only --bus * -l
Path                             Bus
────────────────────────────────────────
/soc/uart@40002000               uart
/soc/i2c@40003000                i2c
/soc/spi@40004000                spi
/soc/usbd@40027000               usb
/soc/uart@40028000               uart
/soc/qspi@40029000               qspi
/soc/qspi@40029000/mx25r6435f@0  on qspi
/soc/spi@4002f000                spi
```

5. Find nodes by interrupts

To find all nodes that generate IRQs:

```
❯ find --interrupt * -l
Path                        Interrupts
──────────────────────────────────────
/soc/clock@40000000         IRQ_0/1
/soc/power@40000000         IRQ_0/1
/soc/radio@40001000         IRQ_1/1
```

To find nodes by interrupt number:

```
/
❯ find --interrupt 28 -l
Path               Interrupts
─────────────────────────────
/soc/pwm@4001c000  IRQ_28/1
```

6. Custom search and *reporting*

To dump all enabled bus devices that generate IRQs,
configuring the table columns with the `-f` option:

```
/
❯ find -c --enabled-only --bus * --interrupt * -f naibcd
Name  Address     Interrupts  Bus   Compatible        Description
───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
uart  0x40002000  IRQ_2/1     uart  nordic,nrf-uarte  Nordic nRF family UARTE (UART with EasyDMA)
i2c   0x40003000  IRQ_3/1     i2c   nordic,nrf-twi    Nordic nRF family TWI (TWI master)…
spi   0x40004000  IRQ_4/1     spi   nordic,nrf-spi    Nordic nRF family SPI (SPI master)
usbd  0x40027000  IRQ_39/1    usb   nordic,nrf-usbd   Nordic nRF52 USB device controller
uart  0x40028000  IRQ_40/1    uart  nordic,nrf-uarte  Nordic nRF family UARTE (UART with EasyDMA)
qspi  0x40029000  IRQ_41/1    qspi  nordic,nrf-qspi   Properties defining the interface for the Nordic QSPI peripheral…
spi   0x4002f000  IRQ_47/1    spi   nordic,nrf-spim   Nordic nRF family SPIM (SPI master with EasyDMA)

7 nodes.
```

To keep this information *at-hand*,
just rely on `dtsh` command output redirection, e.g:

```
/
❯ find -c --enabled-only --bus * --interrupt * -f naibcd > interrupts.txt
```
"""
    # Search criteria.
    _criteria: List[FindCriterion]

    # Nodes matched during the last search.
    _found: List[Node]

    def __init__(self, shell: Dtsh) -> None:
        super().__init__(
            'find',
            'find devicetree nodes',
            True,
            [
                DtshCommandOption('find by name',       None,  'name', 'pattern'),
                DtshCommandOption('find by compatible', None,  'compat', 'pattern'),
                DtshCommandOption('find by bus device', None,  'bus', 'pattern'),
                DtshCommandOption('find by interrupt',  None,  'interrupt', 'pattern'),
                DtshCommandOption('search only enabled nodes', None, 'enabled-only', None),
                DtshCommandOption('print nodes count',  'c',    None, None),
                DtshCommandOption('quiet, only print nodes count',  'q',    None, None),
                DtshCommandFlagLongFmt(),
                DtshCommandArgLongFmt()
            ]
        )
        self._dtsh = shell
        self._criteria = []
        self._found = []

    @property
    def usage(self) -> str:
        """Overrides DtshCommand.usage().
        """
        return super().usage + ' [PATH]'

    @property
    def arg_pattern_name(self) -> Union[str, None]:
        return self.arg_value('--name')

    @property
    def arg_pattern_compat(self) -> Union[str, None]:
        return self.arg_value('--compat')

    @property
    def arg_pattern_bus(self) -> Union[str, None]:
        return self.arg_value('--bus')

    @property
    def arg_pattern_irq(self) -> Union[str, None]:
        return self.arg_value('--interrupt')

    @property
    def with_only_enabled(self) -> bool:
        return self.with_flag('--enabled-only')

    @property
    def with_node_count(self) -> bool:
        return self.with_flag('-c')

    @property
    def with_quiet(self) -> bool:
        return self.with_flag('-q')

    def reset(self) -> None:
        """Overrides DtshCommand.reset().
        """
        super().reset()
        self._found.clear()
        self._criteria.clear()

    def parse_argv(self, argv: List[str]) -> None:
        """Overrides DtshCommand.parse_argv().
        """
        super().parse_argv(argv)
        if len(self._params) > 1:
            raise DtshCommandUsageError(self, 'too many parameters')
        if self.arg_pattern_name:
            self._criteria.append(FindByNameCriterion(self.arg_pattern_name))
        if self.arg_pattern_compat:
            self._criteria.append(FindByCompatCriterion(self.arg_pattern_compat))
        if self.arg_pattern_bus:
            self._criteria.append(FindByBusCriterion(self.arg_pattern_bus))
        if self.arg_pattern_irq:
            self._criteria.append(FindByIrqCriterion(self.arg_pattern_irq))

    def execute(self, vt: DtshVt) -> None:
        """Implements DtshCommand.execute().
        """
        if self._params:
            arg_path = self._dtsh.realpath(self._params[0])
        else:
            arg_path = self._dtsh.pwd

        self._find_nodes(self._dtsh.path2node(arg_path))
        # Like the POSIX find command, does not touch stdout if no match.
        if len(self._found) > 0:

            if self.with_pager:
                vt.pager_enter()

            if not self.with_quiet:
                longfmt = self.arg_longfmt
                if self.with_longfmt and not longfmt:
                   longfmt = self._mk_default_longmt()

                if longfmt:
                    self._write_found_long(vt, longfmt)
                else:
                    self._write_found_default(vt)

            if self.with_node_count:
                vt.write()
                vt.write(f"{len(self._found)} nodes.")

            if self.with_pager:
                vt.pager_exit()

    def autocomplete_param(self, prefix: str) -> Tuple[int, List]:
        """Overrides DtshCommand.autocomplete_param().
        """
        return (DtshAutocomp.MODE_DT_NODE,
                DtshAutocomp.autocomplete_with_nodes(prefix, self._dtsh))

    def _find_nodes(self, root: Node) -> None:
        if self.with_only_enabled and not Dtsh.is_node_enabled(root):
            return
        if self._match_criteria(root):
            self._found.append(root)
        for node in root.children.values():
            self._find_nodes(node)

    def _match_criteria(self, node: Node) -> bool:
        for criterion in self._criteria:
            if not criterion.match(node):
                return False
        return True

    def _write_found_default(self, vt: DtshVt) -> None:
        for node in self._found:
            vt.write(node.path)

    def _write_found_long(self, vt: DtshVt, longfmt: str) -> None:
        ls_table = LsNodeTable(self._dtsh, longfmt)
        for node in self._found:
            ls_table.add_node_row(node)
        vt.write(ls_table.as_view())

    def _mk_default_longmt(self) -> str:
        longfmt = 'p'
        if self.arg_pattern_name:
            longfmt += 'A'
            longfmt += 'L'
        if self.arg_pattern_irq:
            longfmt += 'i'
        if self.arg_pattern_bus:
            longfmt += 'b'
        if self.arg_pattern_compat:
            longfmt += 'c'
            longfmt += 'd'
        return longfmt

class FindByNameCriterion(FindCriterion):
    """Find nodes by names.

    Match nodes by the node-name component (DTSpec 2.2.1).

    If the search pattern does not contain any wildcard,
    defaults to plain text search.

    If the search pattern contains '*' wildcards,
    a regular expression type search is assumed
    where '*' are replaced by the character class for node names
    (see CLASS_NODE_NAME).

    Then:
    - * will match any node-name
    - <prefix>* will match a node-name if it starts with <prefix>
    - *<postfix> will match a node-name if it ends with <postfix>
    - <prefix>*<postfix> will match a node-name if it starts with <prefix>
      and ends with <postfix>
    """

    # Character class for node names (DTSpec 2.2.1).
    CLASS_NODE_NAME = r'[\w,.+\-]*'

    # Pattern for regular expression type search, None for plain text search.
    _re: Union[re.Pattern, None]

    def __init__(self, pattern: str) -> None:
        """Initialize criterion.

        Arguments:
        pattern - the search pattern
        """
        super().__init__(pattern)
        self._re = None
        if pattern.find('*') != -1:
            pattern = pattern.replace('*', FindByNameCriterion.CLASS_NODE_NAME)
            re_expr = f'^{pattern}$'
            self._re = re.compile(re_expr)

    def match(self, node: Node) -> bool:
        name = DtshTui.get_node_nick(node)
        if self._re:
            # RE type search.
            return self._re.match(name) is not None
        # Plain text type search.
        return name.find(self._pattern) != -1


class FindByCompatCriterion(FindCriterion):
    """Find nodes by compatible string.

    Match nodes by the 'compatible' property value (DTSpec 2.3.1).

    If the search pattern does not contain any wildcard,
    defaults to plain text search.

    If the search pattern contains '*' wildcards,
    a regular expression type search is assumed
    where '*' are replaced by the character class for 'compatible' values
    (see CLASS_COMPATIBLE).

    Then:
    - * will match any compatible string
    - <prefix>* will match a compatible if it starts with <prefix>
    - *<postfix> will match a compatible if it ends with <postfix>
    - <prefix>*<postfix> will match a compatible if it starts with <prefix>
      and ends with <postfix>
    """

    # Character class for 'compatible' property (DTSpec 2.3.1).
    CLASS_COMPATIBLE = r'[a-z0-9\-,]*'

    # Pattern for regular expression type search, None for plain text search.
    _re: Union[re.Pattern, None]

    def __init__(self, pattern: str) -> None:
        """Initialize criterion.

        Arguments:
        pattern - the search pattern
        """
        super().__init__(pattern)
        self._re = None
        if pattern.find('*') != -1:
            pattern = pattern.replace('*', FindByCompatCriterion.CLASS_COMPATIBLE)
            re_expr = f'^{pattern}$'
            self._re = re.compile(re_expr)

    def match(self, node: Node) -> bool:
        for compat in node.compats:
            if self._re:
                # RE type search.
                if self._re.match(compat) is not None:
                    return True
            else:
                # Plain text type search.
                if compat.find(self._pattern) != -1:
                    return True
        return False


class FindByBusCriterion(FindCriterion):
    """Find nodes by bus device.

    If the search pattern does not contain any wildcard,
    defaults to plain text search.

    If the search pattern contains '*' wildcards,
    a regular expression type search is assumed
    where '*' are replaced by the character class for bus devices
    (see CLASS_BUS).

    Then:
    - * will match any bus device
    - <prefix>* will match a bus device if it starts with <prefix>
    - *<postfix> will match a bus device if it ends with <postfix>
    - <prefix>*<postfix> will match a bus device if it starts with <prefix>
      and ends with <postfix>
    """

    # Character class for bus name (alphanumeric ? lowercase ?).
    CLASS_BUS = r'[\w]*'

    # Pattern for regular expression type search, None for plain text search.
    _re: Union[re.Pattern, None]

    def __init__(self, pattern: str) -> None:
        """Initialize criterion.

        Arguments:
        pattern - the search pattern
        """
        super().__init__(pattern)
        self._re = None
        if pattern.find('*') != -1:
            pattern = pattern.replace('*', FindByBusCriterion.CLASS_BUS)
            re_expr = f'^{pattern}$'
            self._re = re.compile(re_expr)

    def match(self, node: Node) -> bool:
        for bus in node.buses:
            if self._re:
                if self._re.match(bus) is not None:
                    return True
            else:
                if bus.find(self._pattern) != -1:
                    return True
        for bus in node.on_buses:
            if self._re:
                if self._re.match(bus) is not None:
                    return True
            else:
                if bus.find(self._pattern) != -1:
                    return True
        return False


class FindByIrqCriterion(FindCriterion):
    """Find nodes by interrupts.

    If the search pattern successfully converts to an integer,
    a search by IRQ number is assumed.

    If the conversion fails, a search by interrupt name is assumed.

    If the search pattern does not contain any wildcard,
    defaults to plain text search.

    If the search pattern contains '*' wildcards,
    a regular expression type search is assumed
    where '*' are replaced by the character class for IRQ names
    (see CLASS_IRQ_NAME).

    Then:
    - * will match any IRQ name
    - <prefix>* will match an IRQ name if it starts with <prefix>
    - *<postfix> will match an IRQ name if it ends with <postfix>
    - <prefix>*<postfix> will match an IRQ name if it starts with <prefix>
      and ends with <postfix>
    """

    # Character class for IRQ names (empiric ?).
    CLASS_IRQ_NAME = r'[\w\-]*'

    # Search by IRQ number, None for search by IRQ name.
    _irq_num: Union[int, None]

    # Search by IRQ name: pattern for regular expression type search,
    # None for plain text search.
    _re: Union[re.Pattern, None]

    def __init__(self, pattern: str) -> None:
        """Initialize criterion.

        Arguments:
        pattern - the search pattern
        """
        super().__init__(pattern)
        self._irq_num = None
        self._re = None

        # We do NOT actually configure the criterion filters when
        # the pattern equals to '*': this allows to use find to list all
        # interrupts, including interrupts without name.
        if pattern != '*':
            try:
                self._irq_num = int(pattern)
            except ValueError:
                if pattern.find('*') != -1:
                    pattern = pattern.replace('*', FindByIrqCriterion.CLASS_IRQ_NAME)
                    self._re = re.compile(f'^{pattern}$')

    def match(self, node: Node) -> bool:
        for irq in node.interrupts:
            if self._pattern == '*':
                # Allows find to list all interrupts with '--interrupt *'.
                return True
            if self._irq_num is not None:
                 # Search by IRQ number
                if irq.data.get('irq') == self._irq_num:
                    return True
            else:
                # Search by IRQ name
                if self._re is not None:
                    # RE type search
                    if irq.name and (self._re.match(irq.name) is not None):
                        return True
                else:
                    # Plain text search
                    if irq.name:
                        if irq.name.find(self._pattern) != -1:
                            return True
        return False
