from __future__ import annotations

from pathlib import Path

from cognite.client import CogniteClient


def _load_toml(path: Path) -> dict:
    try:
        import toml
    except ImportError:  # pragma: no cover - python<3.11 path
        import tomli as toml  # type: ignore[no-redef]

    return toml.load(path)


def load_cognite_client_from_toml(
    toml_file: Path | str = "config.toml", section: str | None = "cognite"
) -> CogniteClient:
    """Load a CogniteClient from a TOML file.

    Args:
        toml_file: Path to the TOML file.
        section: Section name to use. If None, use the top-level keys.

    Returns:
        CogniteClient configured with OAuth credentials from TOML.
    """
    path = Path(toml_file)
    if not path.exists():
        raise FileNotFoundError(f"TOML credentials file not found: {path}")

    toml_content = _load_toml(path)
    if section is not None:
        if section not in toml_content:
            raise KeyError(f"Missing section '{section}' in TOML file: {path}")
        toml_content = toml_content[section]

    login_flow = toml_content.pop("login_flow", None)
    if login_flow == "interactive":
        return CogniteClient.default_oauth_interactive(**toml_content)
    return CogniteClient.default_oauth_client_credentials(**toml_content)
