from collections import defaultdict
from typing import Any, Dict, List, NamedTuple, Optional, Set, Tuple, Union

from tomlkit import TOMLDocument, comment, document, inline_table, item, table
from tomlkit.items import InlineTable, Table

from conda_lock._export_lock_spec_compute_platform_indep import (
    EditableDependency,
    unify_platform_independent_deps,
)
from conda_lock.common import warn
from conda_lock.models.lock_spec import (
    Dependency,
    LockSpecification,
    URLDependency,
    VCSDependency,
    VersionedDependency,
)


class TomlTableKey(NamedTuple):
    """Represents a key in a pixi.toml table.

    It can be rendered into a TOML header using `toml_header_sequence`.

    >>> toml_header_sequence(
    ...     TomlTableKey(category="dev", platform="linux-64", manager="pip")
    ... )
    ['feature', 'dev', 'target', 'linux-64', 'pypi-dependencies']
    """

    category: str
    platform: Optional[str]
    manager: str


def render_pixi_toml(
    *,
    lock_spec: LockSpecification,
    project_name: Optional[str] = None,
    with_cuda: Optional[str] = None,
    editables: Optional[List[EditableDependency]] = None,
) -> TOMLDocument:
    """Render a pixi.toml from a LockSpecification as a tomlkit TOMLDocument."""
    pixi_toml = document()
    for line in (
        "This file was generated by conda-lock for the pixi environment manager.",
        "For more information, see <https://github.com/conda/conda-lock> "
        "and <https://pixi.sh>.",
        "Source files:",
        *(f"- {src_file.as_posix()}" for src_file in lock_spec.sources),
    ):
        pixi_toml.add(comment(line))

    if project_name is None:
        project_name = "project-name-placeholder"
    all_platforms = sorted(lock_spec.dependencies.keys())
    all_categories: Set[str] = set()
    for platform in all_platforms:
        for dep in lock_spec.dependencies[platform]:
            all_categories.add(dep.category)
    if {"main", "default"} <= all_categories:
        # The default "main" category of conda-lock will be rendered as
        # the "default" feature in pixi. If both are defined, then this
        # is a conflict.
        raise ValueError("Cannot have both 'main' and 'default' as categories/extras")

    pixi_toml.add(
        "project",
        item(
            dict(
                name=project_name,
                platforms=list(all_platforms),
                channels=[channel.url for channel in lock_spec.channels],
            )
        ),
    )
    if len(lock_spec.channels) == 0:
        warn(
            "No channels defined in the lock file! Consider defining channels, e.g. "
            "via the command line by adding `--channel conda-forge`."
        )
    for channel in lock_spec.channels:
        if channel.used_env_vars:
            warn(
                f"Channel {channel.url} uses environment variables, which will "
                "be dropped in the pixi.toml."
            )

    # The dependency tables
    arranged_deps = arrange_for_toml(lock_spec, editables=editables)
    for key, deps_by_name in arranged_deps.items():
        header_sequence: List[str] = toml_header_sequence(key)

        # Keys are package names, values are version numbers or matchspecs.
        inner_dict = {
            name: toml_dependency_value(dep) for name, dep in deps_by_name.items()
        }

        # Using the nested sequence of headers, walk down the tree of tomlkit tables
        # towards the leaf where the inner_dict should be inserted.
        node: Union[TOMLDocument, Table] = pixi_toml
        for header in header_sequence:
            if header not in node:
                node.add(header, table())
            next_node = node[header]
            assert isinstance(next_node, Table)
            node = next_node
        # Now `node` is the leaf table where the inner table should be inserted.
        node.update(inner_dict)

    # The environments table
    if len(all_categories) > 1:
        pixi_toml.add("environments", toml_environments_table(all_categories))

    # The system requirements table
    if with_cuda:
        pixi_toml.add("system-requirements", item(dict(cuda=with_cuda)))

    return pixi_toml


def toml_dependency_value(
    dep: Union[Dependency, EditableDependency],
) -> Union[str, InlineTable]:
    """Render a conda-lock Dependency as a pixi.toml line as VersionSpec or matchspec.

    The result is suitable for the values used in the `dependencies` or
    `pypi-dependencies` tables of a pixi TOML file.

    >>> toml_dependency_value(VersionedDependency(name="numpy", version="2.1.1"))
    '2.1.1'

    >>> toml_dependency_value(VersionedDependency(name="numpy", version=""))
    '*'

    >>> toml_dependency_value(
    ...     VersionedDependency(
    ...         name="numpy",
    ...         version="2.1.1",
    ...         conda_channel="conda-forge",
    ...         build="py313h4bf6692_0"
    ...     )
    ... )
    {'version': '2.1.1', 'build': 'py313h4bf6692_0', 'channel': 'conda-forge'}

    >>> toml_dependency_value(
    ...     VersionedDependency(
    ...         name="xarray",
    ...         version="",
    ...         extras=["io", "parallel"],
    ...         manager="pip",
    ...     )
    ... )
    {'version': '*', 'extras': ['io', 'parallel']}
    """
    matchspec: Dict[str, Any] = {}
    if isinstance(dep, VersionedDependency):
        matchspec["version"] = dep.version or "*"
        if dep.manager == "pip" and matchspec["version"][0].isdigit():
            matchspec["version"] = f"=={matchspec['version']}"
        if dep.build is not None:
            matchspec["build"] = dep.build
        if dep.hash is not None:
            raise NotImplementedError(f"Hash not yet supported in {dep}")
        if dep.conda_channel is not None:
            matchspec["channel"] = dep.conda_channel
        if dep.extras:
            if dep.extras and dep.manager == "conda":
                warn(f"Extras not supported in Conda dep {dep}")
            else:
                matchspec["extras"] = dep.extras
        if len(matchspec) == 1:
            # Use the simpler VersionSpec format if there's only a version.
            return matchspec["version"]
        else:
            return _dict_to_inline_table(matchspec)
    elif isinstance(dep, URLDependency):
        raise NotImplementedError(f"URL not yet supported in {dep}")
    elif isinstance(dep, VCSDependency):
        raise NotImplementedError(f"VCS not yet supported in {dep}")
    elif isinstance(dep, EditableDependency):
        return _dict_to_inline_table(dict(path=dep.path, editable=True))
    else:
        raise ValueError(f"Unknown dependency type {dep}")


def _dict_to_inline_table(d: Dict[str, Any]) -> InlineTable:
    """Convert a dictionary to a TOML inline table."""
    table = inline_table()
    table.update(d)
    return table


def arrange_for_toml(
    lock_spec: LockSpecification,
    *,
    editables: Optional[List[EditableDependency]] = None,
) -> Dict[TomlTableKey, Dict[str, Union[Dependency, EditableDependency]]]:
    """Arrange dependencies into a structured dictionary for TOML generation."""
    unified_deps = unify_platform_independent_deps(
        lock_spec.dependencies, editables=editables
    )

    # Stick all the dependencies into the correct TOML table
    unsorted_result: Dict[
        TomlTableKey, Dict[str, Union[Dependency, EditableDependency]]
    ] = defaultdict(dict)
    for dep_key, dep in unified_deps.items():
        toml_key = TomlTableKey(
            category=dep_key.category,
            platform=dep_key.platform,
            manager=dep_key.manager,
        )
        if dep_key.name in unsorted_result[toml_key]:
            # This should never happen. Keys must be unique within a TOML table.
            # Moreover, the `unify_platform_independent_deps` function should have
            # already ensured that there are no duplicate keys.
            preexisting_dep = unsorted_result[toml_key][dep_key.name]
            raise RuntimeError(
                f"Duplicate key {dep_key} for {dep} and {preexisting_dep}"
            )
        unsorted_result[toml_key][dep_key.name] = dep

    # Alphabetize the dependencies within each table
    alphabetized_result = {
        toml_key: dict(sorted(deps_by_name.items()))
        for toml_key, deps_by_name in unsorted_result.items()
    }

    # Sort the tables themselves
    sorted_result = dict(sorted(alphabetized_result.items(), key=toml_ordering))

    return sorted_result


def toml_ordering(item: Tuple[TomlTableKey, dict]) -> Tuple[str, str, str]:
    """Make a sort key to properly order the dependency tables in the pixi.toml.

    The main category = default feature comes first. Then the other categories.

    Within each category, the platform-independent dependencies come first, followed by
    the platform-specific dependencies.

    Within each platform, we declare the conda dependencies first, followed by the pip
    dependencies.

    Within each table determined by the hierarchy of category, platform, and manager,
    the dependencies are sorted alphabetically by name. But the key here is just for
    sorting the tables that occur, not the dependencies within them.

    We define the ordering via a tuple of strings: (category, platform, manager).
    The main category and the platform-independent dependencies are represented by
    empty strings so that they come lexicographically first.

    >>> toml_ordering(
    ...     (TomlTableKey(category="main", platform=None, manager="conda"), {})
    ... )
    ('', '', 'conda')

    >>> toml_ordering(
    ...     (TomlTableKey(category="main", platform="linux-64", manager="conda"), {})
    ... )
    ('', 'linux-64', 'conda')

    >>> toml_ordering((TomlTableKey(category="dev", platform=None, manager="pip"), {}))
    ('dev', '', 'pip')
    """
    key = item[0]
    category = "" if key.category in ["main", "default"] else key.category
    platform = key.platform if key.platform is not None else ""
    # "conda" before "pip" is conveniently already lexicographical order.
    return category, platform, key.manager


def toml_header_sequence(key: TomlTableKey) -> List[str]:
    """Generates a TOML header based on the dependency type, platform, and manager.

    >>> toml_header_sequence(
    ...     TomlTableKey(category="main", platform=None, manager="conda")
    ... )
    ['dependencies']

    >>> toml_header_sequence(
    ...     TomlTableKey(category="main", platform="linux-64", manager="conda")
    ... )
    ['target', 'linux-64', 'dependencies']

    >>> toml_header_sequence(
    ...     TomlTableKey(category="main", platform=None, manager="pip")
    ... )
    ['pypi-dependencies']

    >>> toml_header_sequence(
    ...     TomlTableKey(category="main", platform="linux-64", manager="pip")
    ... )
    ['target', 'linux-64', 'pypi-dependencies']

    >>> toml_header_sequence(
    ...     TomlTableKey(category="dev", platform=None, manager="conda")
    ... )
    ['feature', 'dev', 'dependencies']

    >>> toml_header_sequence(
    ...     TomlTableKey(category="dev", platform="linux-64", manager="conda")
    ... )
    ['feature', 'dev', 'target', 'linux-64', 'dependencies']

    >>> toml_header_sequence(TomlTableKey(category="dev", platform=None, manager="pip"))
    ['feature', 'dev', 'pypi-dependencies']

    >>> toml_header_sequence(
    ...     TomlTableKey(category="dev", platform="linux-64", manager="pip")
    ... )
    ['feature', 'dev', 'target', 'linux-64', 'pypi-dependencies']
    """
    parts = []
    if key.category not in ["main", "default"]:
        parts.extend(["feature", key.category])
    if key.platform:
        parts.extend(["target", key.platform])
    parts.append("dependencies" if key.manager == "conda" else "pypi-dependencies")
    return parts


def toml_environments_table(all_categories: Set[str]) -> Table:
    r"""Define the environments section of a pixi.toml file.

    >>> environments_table = toml_environments_table({"main", "dev", "docs"})
    >>> print(environments_table.as_string())
    # Redefine the default environment to include all categories.
    default = ["dev", "docs"]
    # Define a minimal environment with only the default feature.
    minimal = []
    # Create an environment for each feature.
    dev = ["dev"]
    docs = ["docs"]
    <BLANKLINE>
    """
    non_default_categories = sorted(all_categories - {"main", "default"})
    if len(non_default_categories) == 0:
        raise ValueError("Expected at least one non-default category")

    environments = table()
    environments.add(
        comment("Redefine the default environment to include all categories.")
    )
    environments.add("default", non_default_categories)

    MINIMAL_ENVIRONMENT_NAMES = ["minimal", "prod", "main"]
    minimal_category_name = next(
        (name for name in MINIMAL_ENVIRONMENT_NAMES if name not in all_categories), None
    )
    if minimal_category_name is None:
        warn(
            "Can't find a name for the 'minimal' environment since categories for '"
            + "', '".join(MINIMAL_ENVIRONMENT_NAMES)
            + "' are already defined. Skipping."
        )
    else:
        environments.add(
            comment("Define a minimal environment with only the default feature.")
        )
        environments.add(minimal_category_name, [])

    environments.add(comment("Create an environment for each feature."))
    for category in non_default_categories:
        environments.add(category, [category])
    return environments