#!/usr/bin/env python3
#
# Copyright (c) 2022, Nordic Semiconductor ASA
#
# SPDX-License-Identifier: Apache-2.0

'''Internal snippets tool.
This is part of the build system's support for snippets.
It is not meant for use outside of the build system.

Output CMake variables:

- SNIPPET_NAMES: CMake list of discovered snippet names
- SNIPPET_FOUND_{snippet}: one per discovered snippet
'''

from collections import defaultdict, UserDict
from dataclasses import dataclass, field
from pathlib import Path, PurePosixPath
from collections.abc import Iterable
from jsonschema.exceptions import best_match
import argparse
import logging
import os
import re
import sys
import yaml
import platform
import jsonschema

try:
    # Use the C LibYAML parser if available, rather than the Python parser.
    from yaml import CSafeLoader as SafeLoader
except ImportError:
    from yaml import SafeLoader     # type: ignore

# Marker type for an 'append:' configuration. Maps variables
# to the list of values to append to them.
Appends = dict[str, list[str]]
BoardRevisionAppends = dict[str, dict[str, list[str]]]

def _new_append():
    return defaultdict(list)

def _new_board2appends():
    return defaultdict(lambda: defaultdict(_new_append))

def _new_dirs():
    return []

@dataclass
class Snippet:
    '''Class for keeping track of all the settings discovered for an
    individual snippet.'''

    name: str
    dirs: list[Path] = field(default_factory=_new_dirs)
    description: str | None = None
    appends: Appends = field(default_factory=_new_append)
    board2appends: dict[str, BoardRevisionAppends] = field(default_factory=_new_board2appends)

    def process_data(self, pathobj: Path, snippet_data: dict, sysbuild: bool):
        '''Process the data in a snippet.yml file, after it is loaded into a
        python object and validated by jsonschema.'''
        def append_value(variable, value):
            if variable in ('SB_EXTRA_CONF_FILE', 'EXTRA_DTC_OVERLAY_FILE', 'EXTRA_CONF_FILE'):
                paths = []
                if not isinstance(value, list):
                    value = [value]
                for file in value:
                    path = pathobj.parent / file
                    if not path.is_file():
                        _err(f'snippet file {pathobj}: {variable}: file not found: {path}')
                    paths.append(f'"{path.as_posix()}"')
                return paths
            if variable in ('DTS_EXTRA_CPPFLAGS'):
                return [f'"{value}"']
            _err(f'unknown append variable: {variable}')

        for variable, value in snippet_data.get('append', {}).items():
            if (sysbuild is True and variable[0:3] == 'SB_') or \
            (sysbuild is False and variable[0:3] != 'SB_'):
                self.appends[variable] += append_value(variable, value)
        for board, settings in snippet_data.get('boards', {}).items():
            if board.startswith('/') and not board.endswith('/'):
                _err(f"snippet file {pathobj}: board {board} starts with '/', so "
                     "it must end with '/' to use a regular expression")
            for revision, appenddata in settings.get('revisions', {}).items():
                for variable, value in appenddata.get('append', {}).items():
                    if (sysbuild is True and variable[0:3] == 'SB_') or \
                    (sysbuild is False and variable[0:3] != 'SB_'):
                        self.board2appends[board][revision][variable] += \
                            append_value(variable, value)
            for variable, value in settings.get('append', {}).items():
                if (sysbuild is True and variable[0:3] == 'SB_') or \
                (sysbuild is False and variable[0:3] != 'SB_'):
                    self.board2appends[board][""][variable] += \
                        append_value(variable, value)
        self.description = snippet_data.get('description')
        self.dirs.append(pathobj.parent)

class Snippets(UserDict):
    '''Type for all the information we have discovered about all snippets.
    As a dict, this maps a snippet's name onto the Snippet object.
    Any additional global attributes about all snippets go here as
    instance attributes.'''

    def __init__(self, requested: Iterable[str] = None):
        super().__init__()
        self.paths: set[Path] = set()
        self.requested: list[str] = list(requested or [])

class SnippetsError(Exception):
    '''Class for signalling expected errors'''

    def __init__(self, msg):
        self.msg = msg

class SnippetToCMakeOutput:
    '''Helper class for outputting a Snippets's semantics to a .cmake
    include file for use by snippets.cmake.'''

    def __init__(self, snippets: Snippets):
        self.snippets = snippets
        self.section = '#' * 79

    def output_cmake(self):
        '''Output to the file provided to the constructor.'''
        # TODO: add source file info
        snippets = self.snippets
        snippet_names = sorted(snippets.keys())
        output = ''

        if platform.system() == "Windows":
            # Change to linux-style paths for windows to avoid cmake escape character code issues
            snippets.paths = set(map(lambda x: str(PurePosixPath(x)), snippets.paths))

            for this_snippet in snippets:
                for snippet_append in (snippets[this_snippet].appends):
                    snippets[this_snippet].appends[snippet_append] = \
                        set(map(lambda x: str(x.replace("\\", "/")), \
                            snippets[this_snippet].appends[snippet_append]))

        snippet_path_list = " ".join(
            sorted(f'"{path}"' for path in snippets.paths))

        output += '''\
# WARNING. THIS FILE IS AUTO-GENERATED. DO NOT MODIFY!
#
# This file contains build system settings derived from your snippets.
# Its contents are an implementation detail that should not be used outside
# of Zephyr's snippets CMake module.
#
# See the Snippets guide in the Zephyr documentation for more information.
'''

        output += f'''\
{self.section}
# Global information about all snippets.

# The name of every snippet that was discovered.
set(SNIPPET_NAMES {' '.join(f'"{name}"' for name in snippet_names)})
# The paths to all the snippet.yml files. One snippet
# can have multiple snippet.yml files.
set(SNIPPET_PATHS {snippet_path_list})

# Create variable scope for snippets build variables
zephyr_create_scope(snippets)
'''

        for snippet_name in snippets.requested:
            output += self.output_cmake_for(snippets[snippet_name])

        return output

    def output_cmake_for(self, snippet: Snippet):
        output = f'''\
{self.section}
# Snippet '{snippet.name}'

# Common variable appends.
'''
        output += self.output_appends(snippet.appends, 0)
        for board, appends in snippet.board2appends.items():
            output += self.output_appends_for_board(board, appends)
        return output

    def output_appends_for_board(self, board: str, appends: Appends):
        output = ''
        if board.startswith('/'):
            board_re = board[1:-1]
            output += f'''\
# Appends for board regular expression '{board_re}'
if("${{BOARD}}/${{BOARD_QUALIFIERS}}" MATCHES "^{board_re}$")
'''
        else:
            output += f'''\
# Appends for board '{board}'
if("${{BOARD}}/${{BOARD_QUALIFIERS}}" STREQUAL "{board}")
'''

        # Output board variables first then board revision variables
        output += self.output_appends(appends[""], 1)

        for revision in appends:
            if revision != "":
                output += f'''\
  # Appends for revision '{revision}'
  if("${{BOARD_REVISION}}" STREQUAL "{revision}")
'''
                output += self.output_appends(appends[revision], 2)
                output += '  endif()\n'

        output += 'endif()\n'
        return output

    def output_appends(self, appends: Appends, indent: int):
        space = '  ' * indent
        output = ''
        for name, values in appends.items():
            for value in values:
                output += f'{space}zephyr_set({name} {value} SCOPE snippets APPEND)\n'
        return output

# Name of the file containing the jsonschema schema for snippet.yml
# files.
SCHEMA_PATH = str(Path(__file__).parent / 'schemas' / 'snippet-schema.yaml')
with open(SCHEMA_PATH, 'rb') as f:
    SNIPPET_SCHEMA = yaml.load(f.read(), Loader=SafeLoader)

SNIPPET_VALIDATOR_CLASS = jsonschema.validators.validator_for(SNIPPET_SCHEMA)
SNIPPET_VALIDATOR_CLASS.check_schema(SNIPPET_SCHEMA)
SNIPPET_VALIDATOR = SNIPPET_VALIDATOR_CLASS(SNIPPET_SCHEMA)

# The name of the file which contains metadata about the snippets
# being defined in a directory.
SNIPPET_YML = 'snippet.yml'

# Regular expression for validating snippet names. Snippet names must
# begin with an alphanumeric character, and may contain alphanumeric
# characters or underscores. This is intentionally very restrictive to
# keep things consistent and easy to type and remember. We can relax
# this a bit later if needed.
SNIPPET_NAME_RE = re.compile('[A-Za-z0-9][A-Za-z0-9_-]*')

# Logger for this module.
LOG = logging.getLogger('snippets')

def _err(msg):
    raise SnippetsError(f'error: {msg}')

def parse_args():
    parser = argparse.ArgumentParser(description='snippets helper',
                                     allow_abbrev=False)
    parser.add_argument('--snippet-root', default=[], action='append', type=Path,
                        help='''a SNIPPET_ROOT element; may be given
                        multiple times''')
    parser.add_argument('--snippet', dest='snippets', default=[], action='append',
                        help='''a SNIPPET element; may be given
                        multiple times''')
    parser.add_argument('--cmake-out', type=Path,
                        help='''file to write cmake output to; include()
                        this file after calling this script''')
    parser.add_argument('--sysbuild', action="store_true",
                        help='''set if this is running as sysbuild''')
    return parser.parse_args()

def setup_logging():
    logging.basicConfig(level=logging.INFO,
                        format='  %(name)s: %(message)s')

def process_snippets(args: argparse.Namespace) -> Snippets:
    '''Process snippet.yml files under each *snippet_root*
    by recursive search. Return a Snippets object describing
    the results of the search.
    '''
    # This will contain information about all the snippets
    # we discover in each snippet_root element.
    snippets = Snippets(requested=args.snippets)

    # Process each path in snippet_root in order, adjusting
    # snippets as needed for each one.
    for root in args.snippet_root:
        process_snippets_in(root, snippets, args.sysbuild)

    return snippets

def find_snippets_in_roots(requested_snippets, snippet_roots) -> Snippets:
    '''Process snippet.yml files under each *snippet_root*
    by recursive search. Return a Snippets object describing
    the results of the search.
    '''
    # This will contain information about all the snippets
    # we discover in each snippet_root element.
    snippets = Snippets(requested=requested_snippets)

    # Process each path in snippet_root in order, adjusting
    # snippets as needed for each one.
    for root in snippet_roots:
        process_snippets_in(root, snippets, False)

    return snippets

def process_snippets_in(root_dir: Path, snippets: Snippets, sysbuild: bool) -> None:
    '''Process snippet.yml files in *root_dir*,
    updating *snippets* as needed.'''

    if not root_dir.is_dir():
        LOG.warning(f'SNIPPET_ROOT {root_dir} '
                    'is not a directory; ignoring it')
        return

    snippets_dir = root_dir / 'snippets'
    if not snippets_dir.is_dir():
        return

    for dirpath, _, filenames in os.walk(snippets_dir):
        if SNIPPET_YML not in filenames:
            continue

        snippet_yml = Path(dirpath) / SNIPPET_YML
        snippet_data = load_snippet_yml(snippet_yml)
        name = snippet_data['name']
        if name not in snippets:
            snippets[name] = Snippet(name=name)
        snippets[name].process_data(snippet_yml, snippet_data, sysbuild)
        snippets.paths.add(snippet_yml)

def load_snippet_yml(snippet_yml: Path) -> dict:
    '''Load a snippet.yml file *snippet_yml*, validate the contents
    against the schema, and do other basic checks. Return the dict
    of the resulting YAML data.'''

    with open(snippet_yml, 'rb') as f:
        try:
            snippet_data = yaml.load(f.read(), Loader=SafeLoader)
        except yaml.scanner.ScannerError:
            _err(f'snippets file {snippet_yml} is invalid YAML')

    errors = list(SNIPPET_VALIDATOR.iter_errors(snippet_data))

    if errors:
        sys.exit('ERROR: Malformed snippet YAML file: '
                 f'{snippet_yml.as_posix()}\n'
                 f'{best_match(errors).message} in {best_match(errors).json_path}')

    name = snippet_data['name']
    if not SNIPPET_NAME_RE.fullmatch(name):
        _err(f"snippet file {snippet_yml}: invalid snippet name '{name}'; "
             'snippet names must begin with a letter '
             'or number, and may only contain letters, numbers, '
             'dashes (-), and underscores (_)')

    return snippet_data

def check_for_errors(snippets: Snippets) -> None:
    unknown_snippets = sorted(snippet for snippet in snippets.requested
                              if snippet not in snippets)
    if unknown_snippets:
        all_snippets = '\n    '.join(sorted(snippets))
        _err(f'''\
snippets not found: {', '.join(unknown_snippets)}
  Please choose from among the following snippets:
    {all_snippets}''')

def write_cmake_out(snippets: Snippets, cmake_out: Path) -> None:
    '''Write a cmake include file to *cmake_out* which
    reflects the information in *snippets*.

    The contents of this file should be considered an implementation
    detail and are not meant to be used outside of snippets.cmake.'''
    if not cmake_out.parent.exists():
        cmake_out.parent.mkdir()

    snippet_data = SnippetToCMakeOutput(snippets).output_cmake()

    if Path(cmake_out).is_file():
        with open(cmake_out, encoding="utf-8") as fp:
            if fp.read() == snippet_data:
                return

    with open(cmake_out, 'w', encoding="utf-8") as f:
        f.write(snippet_data)

def main():
    args = parse_args()
    setup_logging()
    try:
        snippets = process_snippets(args)
        check_for_errors(snippets)
    except SnippetsError as e:
        LOG.critical(e.msg)
        sys.exit(1)
    write_cmake_out(snippets, args.cmake_out)

if __name__ == "__main__":
    main()