"""Provide common functions shared among different C++ code generation modules."""

import io
import math
import re
from typing import List, Tuple, Optional

from icontract import ensure, require

from aas_core_codegen import intermediate
from aas_core_codegen.common import (
    Stripped,
    assert_never,
    Identifier,
    indent_but_first_line,
)
from aas_core_codegen.cpp import (
    naming as cpp_naming,
)

# See: https://google.github.io/styleguide/cppguide.html#Spaces_vs._Tabs
INDENT = "  "
INDENT2 = INDENT * 2
INDENT3 = INDENT * 3
INDENT4 = INDENT * 4
INDENT5 = INDENT * 5
INDENT6 = INDENT * 6
INDENT7 = INDENT * 7
INDENT8 = INDENT * 8

WARNING = Stripped(
    """\
// This code has been automatically generated by aas-core-codegen.
// Do NOT edit or append."""
)


@ensure(lambda result: result.startswith('L"'))
@ensure(lambda result: result.endswith('"'))
def wstring_literal(text: str) -> Stripped:
    """Generate a C++ string literal from the ``text``."""
    escaped = []  # type: List[str]

    for character in text:
        code_point = ord(character)

        if character == "\a":
            escaped.append("\\a")
        elif character == "\b":
            escaped.append("\\b")
        elif character == "\f":
            escaped.append("\\f")
        elif character == "\n":
            escaped.append("\\n")
        elif character == "\r":
            escaped.append("\\r")
        elif character == "\t":
            escaped.append("\\t")
        elif character == "\v":
            escaped.append("\\v")
        elif character == '"':
            escaped.append('\\"')
        elif character == "\\":
            escaped.append("\\\\")
        elif code_point < 32:
            # Non-printable ASCII characters
            escaped.append(f"\\x{ord(character):x}")
        elif code_point <= 127:
            # ASCII
            escaped.append(character)
        elif 127 < code_point < 255:
            # Above ASCII, but can be encoded as a single byte
            escaped.append(f"\\x{ord(character):x}")
        elif 255 <= code_point < 65536:
            # Above ASCII
            escaped.append(f"\\u{ord(character):04x}")
        elif code_point >= 65536:
            # Above Unicode Binary Multilingual Pane
            escaped.append(f"\\U{ord(character):08x}")
        else:
            escaped.append(character)

    # NOTE (mristin, 2023-06-28):
    # We use std::wstring, therefore ``L`` prefix.
    return Stripped('L"{}"'.format("".join(escaped)))


@require(lambda character: len(character) == 1)
def wchar_literal(character: str) -> Stripped:
    """Generate a C++ wide character literal from the ``character``."""
    code_point = ord(character)

    escaped: str

    if character == "\a":
        escaped = "L'\\a'"
    elif character == "\b":
        escaped = "L'\\b'"
    elif character == "\f":
        escaped = "L'\\f'"
    elif character == "\n":
        escaped = "L'\\n'"
    elif character == "\r":
        escaped = "L'\\r'"
    elif character == "\t":
        escaped = "L'\\t'"
    elif character == "\v":
        escaped = "L'\\v'"
    elif character == "'":
        escaped = "L'\\''"
    elif character == "\\":
        escaped = "L'\\\\'"
    elif code_point < 32:
        # Non-printable ASCII characters
        escaped = f"L'\\x{ord(character):x}'"
    elif code_point <= 127:
        # ASCII
        escaped = f"L'{character}'"
    elif 127 < code_point < 255:
        # Above ASCII, but can be encoded as a single byte
        escaped = f"L'\\x{ord(character):x}'"
    elif 0xD800 <= code_point <= 0xDFFF:
        # NOTE (mristin):
        # These are the surrogate points and can not be represented as wide character
        # literals directly as common compilers such as MSVC++ will complain.
        #
        # We have to fool the compiler at this point as we deliberately want to model
        # the surrogate point.
        escaped = f"static_cast<wchar_t>(0x{ord(character):04x})"

    # NOTE (mristin):
    # Mind the intersecting range for surrogate points just above if you ever convert
    # this if-elif-else statement into a mapping or pattern matching.
    elif 255 <= code_point < 65536:
        # Above ASCII
        escaped = f"L'\\u{ord(character):04x}'"
    elif code_point >= 65536:
        # Above Unicode Basic Multilingual Pane
        escaped = f"L'\\U{ord(character):08x}'"
    else:
        raise AssertionError(f"Unexpected unhandled character: {character!r}")

    return Stripped(escaped)


# fmt: off
# NOTE (mristin, 2023-09-24):
# We use a pre-condition here to simplify the client code. The client must check
# before if the input text is all in ASCII, and report to the user if there are
# any non-ASCII characters in the input.
@require(
    lambda text:
    all(
        ord(character) <= 127
        for character in text
    ),
    "Only ASCII text can be converted to a C++ string literal, otherwise encoding "
    "must be assumed."
)
@ensure(lambda result: result.startswith('"'))
@ensure(lambda result: result.endswith('"'))
# fmt: on
def string_literal(text: str) -> Stripped:
    """Generate a C++ string literal from the ``text``."""
    escaped = []  # type: List[str]

    for character in text:
        code_point = ord(character)

        if character == "\a":
            escaped.append("\\a")
        elif character == "\b":
            escaped.append("\\b")
        elif character == "\f":
            escaped.append("\\f")
        elif character == "\n":
            escaped.append("\\n")
        elif character == "\r":
            escaped.append("\\r")
        elif character == "\t":
            escaped.append("\\t")
        elif character == "\v":
            escaped.append("\\v")
        elif character == '"':
            escaped.append('\\"')
        elif character == "\\":
            escaped.append("\\\\")
        elif code_point < 32:
            # Non-printable ASCII characters
            escaped.append(f"\\x{ord(character):x}")
        elif code_point <= 127:
            escaped.append(character)
        else:
            # Above ASCII
            raise ValueError(r"Unexpected non-ascii code point: {character!r}")

    return Stripped('"{}"'.format("".join(escaped)))


def needs_escaping(text: str) -> bool:
    """Check whether the ``text`` contains a character that needs escaping."""
    for character in text:
        code_point = ord(character)

        if character == "\a":
            return True
        elif character == "\b":
            return True
        elif character == "\f":
            return True
        elif character == "\n":
            return True
        elif character == "\r":
            return True
        elif character == "\t":
            return True
        elif character == "\v":
            return True
        elif character == '"':
            return True
        elif character == "\\":
            return True
        elif code_point < 32:
            # Non-printable ASCII characters
            return True
        elif code_point <= 127:
            # ASCII
            continue
        elif 127 < code_point < 255:
            # Above ASCII, but can be encoded as a single byte
            return True
        elif 255 <= code_point < 65536:
            # Above ASCII
            return True
        elif code_point >= 65536:
            # Above Unicode Binary Multilingual Pane
            return True
        else:
            raise AssertionError(f"Unexpected unhandled character: {character!r}")

    return False


def boolean_literal(value: bool) -> Stripped:
    """Generate the boolean literal corresponding to the ``value``."""
    return Stripped("true") if value else Stripped("false")


def float_literal(value: float) -> Stripped:
    """Generate the float literal.

    We assume that the precision of the literal is not critical and rely on
    Python's ``str(.)`` function. However, if you want to specify the exact
    number, you have to format the number yourself, probably using G17 representation.
    """
    if math.isnan(value):
        return Stripped("std::numeric_limits<double>::quiet_NaN()")
    if value == math.inf:
        return Stripped("std::numeric_limits<double>::infinity()")
    elif value == -math.inf:
        return Stripped("-std::numeric_limits<double>::infinity()")
    else:
        return Stripped(str(value))


def bytes_literal(value: bytes) -> Tuple[Stripped, bool]:
    """
    Generate an expression representing the ``value``.

    If there are more than 8 bytes, a multi-line expression is returned.

    :param value: to be represented
    :return: (C++ expression, is multi-line)
    """
    if len(value) == 0:
        return Stripped("std::vector<std::uint8_t>()"), False

    writer = io.StringIO()

    # noinspection PyUnusedLocal
    multi_line = None  # type: Optional[bool]

    if len(value) <= 8:
        items_joined = ", ".join(f"0x{byte:02x}" for byte in value)
        return Stripped(f"{{{items_joined}}}"), False
    else:
        writer.write(
            """\
{"""
        )

        for start in range(0, len(value), 8):
            if start == 0:
                writer.write(f"\n{INDENT}")
            else:
                writer.write(f",\n{INDENT}")

            end = min(start + 8, len(value))

            assert start < end

            for i, byte in enumerate(value[start:end]):
                if i > 0:
                    writer.write(", ")

                writer.write(f"0x{byte:02x}")

        writer.write("\n}")
        multi_line = True

        return Stripped(writer.getvalue()), multi_line


PRIMITIVE_TYPE_MAP = {
    intermediate.PrimitiveType.BOOL: Stripped("bool"),
    intermediate.PrimitiveType.INT: Stripped("int64_t"),
    intermediate.PrimitiveType.FLOAT: Stripped("double"),
    intermediate.PrimitiveType.STR: Stripped("std::wstring"),
    intermediate.PrimitiveType.BYTEARRAY: Stripped("std::vector<std::uint8_t>"),
}


def _assert_all_primitive_types_are_mapped() -> None:
    """Assert that we have explicitly mapped all the primitive types to C++."""
    all_primitive_literals = set(literal.value for literal in PRIMITIVE_TYPE_MAP)

    mapped_primitive_literals = set(
        literal.value for literal in intermediate.PrimitiveType
    )

    all_diff = all_primitive_literals.difference(mapped_primitive_literals)
    mapped_diff = mapped_primitive_literals.difference(all_primitive_literals)

    messages = []  # type: List[str]
    if len(mapped_diff) > 0:
        messages.append(
            f"More primitive maps are mapped than there were defined "
            f"in the ``intermediate._types``: {sorted(mapped_diff)}"
        )

    if len(all_diff) > 0:
        messages.append(
            f"One or more primitive types in the ``intermediate._types`` were not "
            f"mapped in PRIMITIVE_TYPE_MAP: {sorted(all_diff)}"
        )

    if len(messages) > 0:
        raise AssertionError("\n\n".join(messages))


_assert_all_primitive_types_are_mapped()

COMMON_NAMESPACE = Identifier("common")
TYPES_NAMESPACE = Identifier("types")
CONSTANTS_NAMESPACE = Identifier("constants")
VERIFICATION_NAMESPACE = Identifier("verification")
JSONIZATION_NAMESPACE = Identifier("jsonization")
XMLIZATION_NAMESPACE = Identifier("xmlization")
REVM_NAMESPACE = Identifier("revm")
PATTERN_NAMESPACE = Identifier("pattern")


def generate_primitive_type(primitive_type: intermediate.PrimitiveType) -> Stripped:
    """Generate the C++ type for the given primitive type."""
    return PRIMITIVE_TYPE_MAP[primitive_type]


def primitive_type_is_referencable(primitive_type: intermediate.PrimitiveType) -> bool:
    """Return ``True`` if the primitive type denotes a referencable value in C++."""
    if primitive_type is intermediate.PrimitiveType.BOOL:
        return False

    elif primitive_type is intermediate.PrimitiveType.INT:
        return False

    elif primitive_type is intermediate.PrimitiveType.FLOAT:
        return False

    elif primitive_type is intermediate.PrimitiveType.STR:
        return True

    elif primitive_type is intermediate.PrimitiveType.BYTEARRAY:
        return True
    else:
        assert_never(primitive_type)


def generate_primitive_type_with_const_ref_if_applicable(
    primitive_type: intermediate.PrimitiveType,
) -> Stripped:
    """Generate the C++ type and wrap it in ``const T&``, if applicable."""
    code = generate_primitive_type(primitive_type)

    if primitive_type_is_referencable(primitive_type):
        return Stripped(f"const {code}&")

    return code


def generate_type(
    type_annotation: intermediate.TypeAnnotationUnion,
    types_namespace: Optional[Identifier] = None,
    common_namespace: Optional[Identifier] = COMMON_NAMESPACE,
) -> Stripped:
    """
    Generate the C++ type for the given type annotation.

    If ``types_namespace`` is specified, it is prepended to all our types.

    If `common_namespace` is specified, it is prepended to ``optional``.
    """
    if isinstance(type_annotation, intermediate.PrimitiveTypeAnnotation):
        return generate_primitive_type(type_annotation.a_type)

    elif isinstance(type_annotation, intermediate.OurTypeAnnotation):
        our_type = type_annotation.our_type

        if isinstance(our_type, intermediate.Enumeration):
            enum_name = cpp_naming.enum_name(type_annotation.our_type.name)
            if types_namespace is None:
                return enum_name

            return Stripped(f"{types_namespace}::{enum_name}")

        elif isinstance(our_type, intermediate.ConstrainedPrimitive):
            return generate_primitive_type(our_type.constrainee)

        elif isinstance(
            our_type, (intermediate.AbstractClass, intermediate.ConcreteClass)
        ):
            # NOTE (mristin, 2023-06-28):
            # We always refer to interfaces even in cases of concrete classes without
            # concrete descendants since we want to allow enhancing.
            interface_name = cpp_naming.interface_name(our_type.name)

            type_identifier: str

            if types_namespace is None:
                type_identifier = interface_name
            else:
                type_identifier = f"{types_namespace}::{interface_name}"

            assert not type_identifier.endswith(">")
            return Stripped(f"std::shared_ptr<{type_identifier}>")

    elif isinstance(type_annotation, intermediate.ListTypeAnnotation):
        item_type = generate_type(
            type_annotation=type_annotation.items, types_namespace=types_namespace
        )

        if "<" not in item_type:
            return Stripped(f"std::vector<{item_type}>")

        return Stripped(
            f"""\
std::vector<
{INDENT}{indent_but_first_line(item_type, INDENT)}
>"""
        )

    elif isinstance(type_annotation, intermediate.OptionalTypeAnnotation):
        value_type = generate_type(
            type_annotation=type_annotation.value, types_namespace=types_namespace
        )

        common_namespace_prefix = (
            "" if common_namespace is None else f"{common_namespace}::"
        )

        if "<" not in value_type:
            return Stripped(f"{common_namespace_prefix}optional<{value_type}>")

        return Stripped(
            f"""\
{common_namespace_prefix}optional<
{INDENT}{indent_but_first_line(value_type, INDENT)}
>"""
        )

    else:
        assert_never(type_annotation)

    raise AssertionError("Should not have gotten here")


def is_referencable(type_annotation: intermediate.TypeAnnotationUnion) -> bool:
    """Return ``True`` if the type annotation denotes a referencable value in C++."""
    if isinstance(type_annotation, intermediate.OptionalTypeAnnotation):
        return True

    primitive_type = intermediate.try_primitive_type(type_annotation)
    if primitive_type is not None:
        return primitive_type_is_referencable(primitive_type)

    if isinstance(type_annotation, intermediate.OurTypeAnnotation) and isinstance(
        type_annotation.our_type, intermediate.Enumeration
    ):
        return False

    return True


def generate_type_with_const_ref_if_applicable(
    type_annotation: intermediate.TypeAnnotationUnion,
    types_namespace: Optional[Identifier] = None,
) -> Stripped:
    """
    Generate the C++ type and wrap it in ``const T&``, if applicable.

    If ``types_namespace`` is specified, it is prepended to all our types.
    """
    code = generate_type(
        type_annotation=type_annotation, types_namespace=types_namespace
    )
    if is_referencable(type_annotation):
        return Stripped(f"const {code}&")

    return code


def generate_type_with_ref(
    type_annotation: intermediate.TypeAnnotationUnion,
    types_namespace: Optional[Identifier] = None,
) -> Stripped:
    """
    Generate the C++ type and wrap it as a reference in ``T&``.

    If ``types_namespace`` is specified, it is prepended to all our types.
    """
    code = generate_type(
        type_annotation=type_annotation, types_namespace=types_namespace
    )
    return Stripped(f"{code}&")


_ANGLE_BRACKETS_IN_TYPE_RE = re.compile(r"\s*([<>])\s*")
_INDENT_LIST = [
    INDENT,
    INDENT2,
    INDENT3,
    INDENT4,
    INDENT5,
    INDENT6,
    INDENT7,
    INDENT8,
]


def break_type_in_lines(text: str) -> str:
    """Break the given C++ type in multiple lines, if applicable."""
    if len(text) == 0:
        return ""

    buffer = io.StringIO()

    level = 0
    last_match = None  # type: Optional[re.Match[str]]
    for match in _ANGLE_BRACKETS_IN_TYPE_RE.finditer(text):
        if last_match is None:
            buffer.write(text[: match.start()])
        else:
            buffer.write(text[last_match.end() : match.start()])

        assert (
            0 <= level < len(_INDENT_LIST)
        ), f"Expected at most {len(_INDENT_LIST)} levels, but got: {level=}"

        angle_bracket = match.group(1)
        if angle_bracket == "<":
            buffer.write("<\n" + _INDENT_LIST[level])
            level += 1
        elif angle_bracket == ">":
            level -= 1
            buffer.write(" >")
        else:
            raise AssertionError(f"Unhandled angle bracket: {angle_bracket=}")

        last_match = match

    if last_match is None:
        # NOTE (mristin, 2024-01-18):
        # Not a signle angle bracket has been detected.
        return text

    buffer.write(text[last_match.end() :])

    return buffer.getvalue()


class GeneratorForLoopVariables:
    """
    Generate a unique variable name based on ``item`` stem.

    >>> generator = GeneratorForLoopVariables()

    >>> next(generator)
    'v'

    >>> next(generator)
    'v1'

    >>> next(generator)
    'v2'
    """

    def __init__(self) -> None:
        """Initialize with the zero counter."""
        self.counter = 0

    def __next__(self) -> Identifier:
        """Generate the next variable name."""
        if self.counter == 0:
            result = Identifier("v")
        else:
            result = Identifier(f"v{self.counter}")

        self.counter += 1

        return result


def generate_include_prefix_path(library_namespace: Stripped) -> Stripped:
    """
    Generate the prefix path for the includes.

    >>> generate_include_prefix_path(Stripped("some::name::space"))
    'some/name/space'
    """
    return Stripped(library_namespace.replace("::", "/"))


def generate_namespace_opening(library_namespace: Stripped) -> Stripped:
    """
    Generate the code to open the nested library namespace.

    >>> generate_namespace_opening(Stripped("some::name::space"))
    'namespace some {\\nnamespace name {\\nnamespace space {'
    """
    namespace_parts = library_namespace.split("::")
    return Stripped("\n".join(f"namespace {part} {{" for part in namespace_parts))


def generate_namespace_closing(library_namespace: Stripped) -> Stripped:
    """
    Generate the code to close the nested library namespace.

    >>> generate_namespace_closing(Stripped("some::name::space"))
    '}  // namespace space\\n}  // namespace name\\n}  // namespace some'
    """
    namespace_parts = library_namespace.split("::")
    return Stripped(
        "\n".join(f"}}  // namespace {part}" for part in reversed(namespace_parts))
    )


def include_guard_var(namespace: Stripped) -> Stripped:
    """
    Generate the variable name of the include guard.

    >>> include_guard_var(Stripped("aascore3::types"))
    'AASCORE3_TYPES_GUARD_'
    """
    return Stripped("_".join(namespace.split("::")).upper() + "_GUARD_")


def non_documentation_comment(text: str) -> str:
    r"""
    Generate the non-documentation comment prefixed with ``//``.

    >>> non_documentation_comment('test\nme')
    '// test\n// me'
    """
    lines = []  # type: List[str]
    for line in text.splitlines():
        stripped = line.strip()
        if len(stripped) > 0:
            lines.append(f"// {line}")
        else:
            lines.append("//")

    return "\n".join(lines)
