from typing import Union

from pipenv.vendor.plette.models import Package, PackageCollection
from pipenv.vendor.tomlkit.container import Container, OutOfOrderTableProxy
from pipenv.vendor.tomlkit.items import AoT, Array, Bool, InlineTable, Item, String, Table
from pipenv.vendor.tomlkit.toml_document import TOMLDocument

try:
    import tomllib as toml
except ImportError:
    from pipenv.vendor import tomli as toml

from pipenv.vendor import tomlkit

TOML_DICT_TYPES = Union[Container, Package, PackageCollection, Table, InlineTable]
TOML_DICT_OBJECTS = (Container, Package, Table, InlineTable, PackageCollection)
TOML_DICT_NAMES = [o.__class__.__name__ for o in TOML_DICT_OBJECTS]


def cleanup_toml(tml):
    # Remove all empty lines from TOML.
    toml = "\n".join(line for line in tml.split("\n") if line.strip())
    new_toml = []
    # Add newlines between TOML sections.
    for i, line in enumerate(toml.split("\n")):
        # Skip the first line.
        if line.startswith("[") and i > 0:
            # Insert a newline before the heading.
            new_toml.append("")
        new_toml.append(line)
    # adding new line at the end of the TOML file
    new_toml.append("")
    toml = "\n".join(new_toml)
    return toml


def convert_toml_outline_tables(parsed: TOMLDocument, project) -> TOMLDocument:
    """Converts all outline tables to inline tables."""

    def convert_tomlkit_table(section):
        result: Table = tomlkit.table()
        if isinstance(section, Table):
            body = section.value._body
        elif isinstance(section, OutOfOrderTableProxy):
            body = section._internal_container._body
        else:
            assert not hasattr(section, "_body")
            body = section

        index: int = 0
        for key, value in body:
            if not key:
                continue
            if hasattr(value, "keys") and not isinstance(value, InlineTable):
                table = tomlkit.inline_table()
                table.update(value.value)
                key.sep = " = "  # add separator because it did not exist before
                result.append(key, table)
            else:
                result.append(key, value)
            index += 1

        return result

    def convert_toml_table(section):
        result = section.copy()
        for package, value in section.items():
            if hasattr(value, "keys") and not isinstance(
                value, toml.decoder.InlineTableDict
            ):
                table = toml.TomlDecoder().get_empty_inline_table()
                table.update(value)
                result[package] = table
        return result

    is_tomlkit_parsed = isinstance(parsed, Container)
    for section in project.get_package_categories():
        table_data = parsed.get(section, {})
        if table_data is None:
            continue
        if is_tomlkit_parsed:
            result = convert_tomlkit_table(table_data)
        else:
            result = convert_toml_table(table_data)

        parsed[section] = result
    return parsed


def tomlkit_value_to_python(toml_value):
    # type: (Union[Array, AoT, TOML_DICT_TYPES, Item]) -> Union[List, Dict]
    value_type = type(toml_value).__name__
    if (
        isinstance(toml_value, TOML_DICT_OBJECTS + (dict,))
        or value_type in TOML_DICT_NAMES
    ):
        return tomlkit_dict_to_python(toml_value)
    elif isinstance(toml_value, AoT) or value_type == "AoT":
        return [tomlkit_value_to_python(val) for val in toml_value._body]
    elif isinstance(toml_value, Array) or value_type == "Array":
        return [tomlkit_value_to_python(val) for val in list(toml_value)]
    elif isinstance(toml_value, String) or value_type == "String":
        return f"{toml_value!s}"
    elif isinstance(toml_value, Bool) or value_type == "Bool":
        return toml_value.value
    elif isinstance(toml_value, Item):
        return toml_value.value
    return toml_value


def tomlkit_dict_to_python(toml_dict):
    # type: (TOML_DICT_TYPES) -> Dict
    value_type = type(toml_dict).__name__
    if toml_dict is None:
        raise TypeError("Invalid type NoneType when converting toml dict to python")
    converted = None  # type: Optional[Dict]
    if isinstance(toml_dict, (InlineTable, Table)) or value_type in (
        "InlineTable",
        "Table",
    ):
        converted = toml_dict.value
    elif isinstance(toml_dict, (Package, PackageCollection)) or value_type in (
        "Package, PackageCollection"
    ):
        converted = toml_dict._data
        if isinstance(converted, Container) or type(converted).__name__ == "Container":
            converted = converted.value
    elif isinstance(toml_dict, Container) or value_type == "Container":
        converted = toml_dict.value
    elif isinstance(toml_dict, dict):
        converted = toml_dict.copy()
    else:
        raise TypeError(
            f"Invalid type for conversion: expected Container, Dict, or Table, got {toml_dict}"
        )
    if isinstance(converted, dict):
        return {k: tomlkit_value_to_python(v) for k, v in converted.items()}
    elif isinstance(converted, (TOML_DICT_OBJECTS)) or value_type in TOML_DICT_NAMES:
        return tomlkit_dict_to_python(converted)
    return converted
