#!/usr/bin/env python3

# Copyright (c) 2020 Nordic Semiconductor ASA
# SPDX-License-Identifier: Apache-2.0

import argparse
from collections import defaultdict
import itertools
from pathlib import Path
from typing import NamedTuple

ZEPHYR_BASE = Path(__file__).resolve().parents[1]

#
# This is shared code between the build system's 'boards' target
# and the 'west boards' extension command. If you change it, make
# sure to test both ways it can be used.
#
# (It's done this way to keep west optional, making it possible to run
# 'ninja boards' in a build directory without west installed.)
#

class Board(NamedTuple):
    name: str
    arch: str
    dir: Path

def board_key(board):
    return board.name

def find_arch2boards(args):
    arch2board_set = find_arch2board_set(args)
    return {arch: sorted(arch2board_set[arch], key=board_key)
            for arch in arch2board_set}

def find_boards(args):
    return sorted(itertools.chain(*find_arch2board_set(args).values()),
                  key=board_key)

def find_arch2board_set(args):
    arches = sorted(find_arches(args))
    ret = defaultdict(set)

    for root in itertools.chain([ZEPHYR_BASE], args.board_roots):
        for arch, boards in find_arch2board_set_in(root, arches).items():
            ret[arch] |= boards

    return ret

def find_arches(args):
    arch_set = find_arches_in(ZEPHYR_BASE)

    for root in args.arch_roots:
        arch_set |= find_arches_in(root)

    return arch_set

def find_arches_in(root):
    ret = set()
    arch = root / 'arch'
    common = arch / 'common'

    if not arch.is_dir():
        return ret

    for maybe_arch in arch.iterdir():
        if not maybe_arch.is_dir() or maybe_arch == common:
            continue
        ret.add(maybe_arch.name)

    return ret

def find_arch2board_set_in(root, arches):
    ret = defaultdict(set)
    boards = root / 'boards'

    for arch in arches:
        if not (boards / arch).is_dir():
            continue

        for maybe_board in (boards / arch).iterdir():
            if not maybe_board.is_dir():
                continue
            for maybe_defconfig in maybe_board.iterdir():
                file_name = maybe_defconfig.name
                if file_name.endswith('_defconfig'):
                    board_name = file_name[:-len('_defconfig')]
                    ret[arch].add(Board(board_name, arch, maybe_board))

    return ret

def parse_args():
    parser = argparse.ArgumentParser()
    add_args(parser)
    return parser.parse_args()

def add_args(parser):
    # Remember to update west-completion.bash if you add or remove
    # flags
    parser.add_argument("--arch-root", dest='arch_roots', default=[],
                        type=Path, action='append',
                        help='''add an architecture root (ZEPHYR_BASE is
                        always present), may be given more than once''')
    parser.add_argument("--board-root", dest='board_roots', default=[],
                        type=Path, action='append',
                        help='''add a board root (ZEPHYR_BASE is always
                        present), may be given more than once''')

def dump_boards(arch2boards):
    for arch, boards in arch2boards.items():
        print(f'{arch}:')
        for board in boards:
            print(f'  {board.name}')

if __name__ == '__main__':
    dump_boards(find_arch2boards(parse_args()))