# Copyright 2020 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Tools for generating Pigweed tests that execute in C++ and Python."""

import argparse
from dataclasses import dataclass
from datetime import datetime
from collections import defaultdict
import unittest

from typing import (Any, Callable, Dict, Generic, Iterable, Iterator, List,
                    Sequence, TextIO, TypeVar, Union)

_CPP_HEADER = f"""\
// Copyright {datetime.now().year} The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

// AUTOGENERATED - DO NOT EDIT
//
// Generated at {datetime.now().isoformat()}

// clang-format off
"""


class Error(Exception):
    """Something went wrong when generating tests."""


T = TypeVar('T')


@dataclass
class Context(Generic[T]):
    """Info passed into test generator functions for each test case."""
    group: str
    count: int
    total: int
    test_case: T

    def cc_name(self) -> str:
        name = ''.join(w.capitalize()
                       for w in self.group.replace('-', ' ').split(' '))
        name = ''.join(c if c.isalnum() else '_' for c in name)
        return f'{name}_{self.count}' if self.total > 1 else name

    def py_name(self) -> str:
        name = 'test_' + ''.join(c if c.isalnum() else '_'
                                 for c in self.group.lower())
        return f'{name}_{self.count}' if self.total > 1 else name


# Test cases are specified as a sequence of strings or test case instances. The
# strings are used to separate the tests into named groups. For example:
#
#   STR_SPLIT_TEST_CASES = (
#     'Empty input',
#     MyTestCase('', '', []),
#     MyTestCase('', 'foo', []),
#     'Split on single character',
#     MyTestCase('abcde', 'c', ['ab', 'de']),
#     ...
#   )
#
GroupOrTest = Union[str, T]

# Python tests are generated by a function that returns a function usable as a
# unittest.TestCase method.
PyTest = Callable[[unittest.TestCase], None]
PyTestGenerator = Callable[[Context[T]], PyTest]

# C++ tests are generated with a function that returns or yields lines of C++
# code for the given test case.
CcTestGenerator = Callable[[Context[T]], Iterable[str]]


class TestGenerator(Generic[T]):
    """Generates tests for multiple languages from a series of test cases."""
    def __init__(self, test_cases: Sequence[GroupOrTest[T]]):
        self._cases: Dict[str, List[T]] = defaultdict(list)
        message = ''

        if len(test_cases) < 2:
            raise Error('At least one test case must be provided')

        if not isinstance(test_cases[0], str):
            raise Error(
                'The first item in the test cases must be a group name string')

        for case in test_cases:
            if isinstance(case, str):
                message = case
            else:
                self._cases[message].append(case)

        if '' in self._cases:
            raise Error('Empty test group names are not permitted')

    def _test_contexts(self) -> Iterator[Context[T]]:
        for group, test_list in self._cases.items():
            for i, test_case in enumerate(test_list, 1):
                yield Context(group, i, len(test_list), test_case)

    def _generate_python_tests(self, define_py_test: PyTestGenerator):
        tests: Dict[str, Callable[[Any], None]] = {}

        for ctx in self._test_contexts():
            test = define_py_test(ctx)
            test.__name__ = ctx.py_name()

            if test.__name__ in tests:
                raise Error(
                    f'Multiple Python tests are named {test.__name__}!')

            tests[test.__name__] = test

        return tests

    def python_tests(self, name: str, define_py_test: PyTestGenerator) -> type:
        """Returns a Python unittest.TestCase class with tests for each case."""
        return type(name, (unittest.TestCase, ),
                    self._generate_python_tests(define_py_test))

    def _generate_cc_tests(self, define_cpp_test: CcTestGenerator, header: str,
                           footer: str) -> Iterator[str]:
        yield _CPP_HEADER
        yield header

        for ctx in self._test_contexts():
            yield from define_cpp_test(ctx)
            yield ''

        yield footer

    def cc_tests(self, output: TextIO, define_cpp_test: CcTestGenerator,
                 header: str, footer: str):
        """Writes C++ unit tests for each test case to the given file."""
        for line in self._generate_cc_tests(define_cpp_test, header, footer):
            output.write(line)
            output.write('\n')


def _to_chars(data: bytes) -> Iterator[str]:
    for i, byte in enumerate(data):
        try:
            char = data[i:i + 1].decode()
            yield char if char.isprintable() else fr'\x{byte:02x}'
        except UnicodeDecodeError:
            yield fr'\x{byte:02x}'


def cc_string(data: Union[str, bytes]) -> str:
    """Returns a C++ string literal version of a byte string or UTF-8 string."""
    if isinstance(data, str):
        data = data.encode()

    return '"' + ''.join(_to_chars(data)) + '"'


def parse_test_generation_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description='Generate unit test files')
    parser.add_argument('--generate-cc-test',
                        type=argparse.FileType('w'),
                        help='Generate the C++ test file')
    return parser.parse_known_args()[0]