#!/usr/bin/env python3
#
# Copyright (c) 2025 YunoHost Contributors
#
# This file is part of YunoHost (see https://yunohost.org)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

import atexit
import logging
import os
import time
from functools import cache
from typing import Callable, Literal, Mapping

import ldap
import ldap.modlist as modlist
import ldap.sasl
from ldap.ldapobject import LDAPObject, ReconnectLDAPObject
from moulinette import m18n
from moulinette.core import MoulinetteError

from ..utils.error import YunohostError

logger = logging.getLogger("yunohost.utils.ldap")


# We cache the interface to do some caching
# to avoid re-authenticating in case we call _get_ldap_authenticator multiple times
@cache
def _get_ldap_interface() -> "LDAPInterface":
    return LDAPInterface()


# Add this to properly close / delete the ldap interface / authenticator
# when Python exits ...
# Otherwise there's a risk that some funky error appears at the very end
# of the command due to Python stuff being unallocated in wrong order.
def _destroy_ldap_interface() -> None:
    _get_ldap_interface.cache_clear()


atexit.register(_destroy_ldap_interface)


# We regularly want to extract stuff like 'bar' in ldap path like
# foo=bar,dn=users.example.org,ou=example.org,dc=org so this small helper allow
# to do this without relying of dozens of mysterious string.split()[0]
#
# e.g. using _ldap_path_extract(path, "foo") on the previous example will
# return bar
def _ldap_path_extract(path: str, info: str) -> str:
    for element in path.split(","):
        if element.startswith(f"{info}="):
            return element[len(f"{info}=") :]
    raise YunohostError(f"Didnt find info {info} in {path} !?")


def modifyModlist_finegrained(
    old_entry: Mapping[str, str | list[str] | set[str]],
    new_entry: Mapping[str, str | list[str] | set[str]],
) -> list[tuple[int, str, list[str]]]:
    """
    Prepare an optimized modification list to give to ldap.modify_ext()
    """
    ldif = []
    for attribute, value in new_entry.items():
        if not isinstance(value, (list, set)):
            value = {value}
        old_value: str | list[str] | set[str] = old_entry.get(attribute, set())
        if not isinstance(old_value, (list, set)):
            old_value = {old_value.strip()}
        if isinstance(value, set) and value == set(old_value):
            continue
        if isinstance(value, list) and value == list(old_value):
            continue

        if not old_value:
            ldif.append((ldap.MOD_ADD, attribute, list(value)))
        elif not value:
            ldif.append((ldap.MOD_DELETE, attribute, list(old_value)))
        # Add or/and delete only needed values with unordered set
        elif isinstance(value, set):
            values_to_del = set(old_value) - value
            if values_to_del == set(old_value):
                ldif.append((ldap.MOD_REPLACE, attribute, list(value)))
                continue
            elif values_to_del:
                ldif.append((ldap.MOD_DELETE, attribute, list(values_to_del)))

            values_to_add = value - set(old_value)
            if values_to_add:
                ldif.append((ldap.MOD_ADD, attribute, list(values_to_add)))

        # Add or/and delete only needed values with ordered list
        else:
            for i, v in enumerate(value):
                if i >= len(old_value) or old_value[i] != v:  # type: ignore
                    break
            if i == 0:
                ldif.append((ldap.MOD_REPLACE, attribute, value))
            else:
                if old_value[i:]:  # type: ignore
                    ldif.append((ldap.MOD_DELETE, attribute, old_value[i:]))  # type: ignore
                if value[i:]:
                    ldif.append((ldap.MOD_ADD, attribute, value[i:]))

    return ldif


class LDAPInterface:
    URI = "ldapi://%2Fvar%2Frun%2Fslapd%2Fldapi"
    BASEDN = "dc=yunohost,dc=org"
    ROOTDN = "gidNumber=0+uidNumber=0,cn=peercred,cn=external,cn=auth"
    USERDN = "uid={username},ou=users,dc=yunohost,dc=org"

    def __init__(self, user: str | None = None, password: str | None = None) -> None:
        self.userdn: str
        self._connect: Callable[[LDAPObject], None]

        if user is None:
            if os.getuid() == 0:
                logger.debug("initializing root ldap interface")
                self.userdn = self.ROOTDN
                self._connect = lambda con: con.sasl_non_interactive_bind_s("EXTERNAL")
            else:
                logger.debug("initializing anonymous ldap interface")
                self.userdn = ""
                self._connect = lambda con: None
        else:
            logger.debug("initializing user ldap interface")
            self.userdn = self.USERDN.format(username=user)
            self._connect = lambda con: con.simple_bind_s(self.userdn, password)

        self.connect()

    def connect(self) -> None:
        def _reconnect() -> LDAPObject:
            con = ReconnectLDAPObject(self.URI, retry_max=10, retry_delay=2)
            self._connect(con)
            return con

        try:
            con = _reconnect()
        except ldap.SERVER_DOWN:
            # ldap is down, attempt to restart it before really failing
            logger.warning(m18n.n("ldap_server_is_down_restart_it"))
            os.system("systemctl restart slapd")
            time.sleep(10)  # waits 10 secondes so we are sure that slapd has restarted
            try:
                con = _reconnect()
            except ldap.SERVER_DOWN:
                raise YunohostError(
                    "Service slapd is not running but is required to perform this action ... "
                    "You can try to investigate what's happening with 'systemctl status slapd'",
                    raw_msg=True,
                )

        # Check that we are indeed logged in with the right identity
        try:
            # whoami_s return dn:..., then delete these 3 characters
            who = con.whoami_s()[3:]
        except Exception as e:
            logger.warning("Error during ldap authentication process: %s", e)
            raise
        else:
            if who != self.userdn:
                raise MoulinetteError("Not logged in with the expected userdn ?!")
            else:
                self.con = con

    def __del__(self) -> None:
        """Disconnect and free ressources"""
        if hasattr(self, "con") and self.con:
            self.con.unbind_s()

    def search(
        self,
        base: str | None = None,
        filter: str = "(objectClass=*)",
        attrs: set[str] | list[str] | None = ["dn"],
    ) -> list[dict[str, list[str]]]:
        """Search in LDAP base

        Perform an LDAP search operation with given arguments and return
        results as a list.

        Keyword arguments:
            - base -- The dn to search into
            - filter -- A string representation of the filter to apply
            - attrs -- A list of attributes to fetch

        Returns:
            A list of all results

        """
        if not base:
            base = self.BASEDN
        else:
            base = f"{base},{self.BASEDN}"

        try:
            result = self.con.search_s(base, ldap.SCOPE_SUBTREE, filter, attrs)
        except ldap.SERVER_DOWN as e:
            raise e
        except Exception as e:
            raise MoulinetteError(
                "error during LDAP search operation with: base='%s', "
                "filter='%s', attrs=%s and exception %s" % (base, filter, attrs, e),
                raw_msg=True,
            )

        result_list: list[dict[str, list[str]]] = []
        if not attrs or "dn" not in attrs:
            result_list = [entry for dn, entry in result]
        else:
            for dn, entry in result:
                entry["dn"] = [dn]
                result_list.append(entry)

        def decode(value: str | bytes) -> str:
            if isinstance(value, bytes):
                value = value.decode("utf-8")
            return value

        # result_list is for example :
        # [{'virtualdomain': [b'test.com']}, {'virtualdomain': [b'yolo.test']},
        for stuff in result_list:
            if isinstance(stuff, dict):
                for key, values in stuff.items():
                    stuff[key] = [decode(value) for value in values]

        return result_list

    def add(
        self, rdn: str, attr_dict: Mapping[str, str | list[str] | set[str]]
    ) -> bool:
        """
        Add LDAP entry

        Keyword arguments:
            rdn         -- DN without domain
            attr_dict   -- Dictionnary of attributes/values to add

        Returns:
            Boolean | MoulinetteError

        """
        dn = f"{rdn},{self.BASEDN}"
        ldif = modlist.addModlist(attr_dict)
        for i, (k, v) in enumerate(ldif):
            if isinstance(v, list):
                v = [a.encode("utf-8") for a in v]
            elif isinstance(v, str):
                v = [v.encode("utf-8")]
            ldif[i] = (k, v)

        try:
            self.con.add_s(dn, ldif)
        except Exception as e:
            raise MoulinetteError(
                "error during LDAP add operation with: rdn='%s', "
                "attr_dict=%s and exception %s" % (rdn, attr_dict, e),
                raw_msg=True,
            )
        else:
            return True

    def remove(self, rdn: str) -> bool:
        """
        Remove LDAP entry

        Keyword arguments:
            rdn         -- DN without domain

        Returns:
            Boolean | MoulinetteError

        """
        dn = f"{rdn},{self.BASEDN}"
        try:
            self.con.delete_s(dn)
        except Exception as e:
            raise MoulinetteError(
                "error during LDAP delete operation with: rdn='%s' and exception %s"
                % (rdn, e),
                raw_msg=True,
            )
        else:
            return True

    def update(
        self,
        rdn: str,
        attr_dict: Mapping[str, str | list[str] | set[str]],
        new_rdn: str | Literal[False] = False,
    ) -> bool:
        """
        Modify LDAP entry

        Keyword arguments:
            rdn         -- DN without domain
            attr_dict   -- Dictionnary of attributes/values to add
            new_rdn     -- New RDN for modification

        Returns:
            Boolean | MoulinetteError

        """
        dn = f"{rdn},{self.BASEDN}"
        current_entry = self.search(rdn, attrs=None)

        # Previously, we used modifyModlist, which directly uses the lib system libldap
        # supplied with openldap. Unfortunately, the output of this command was not
        # optimal with attributes containing lists (complete deletion then complete
        # rewriting of the list). In view of the major performance problems associated
        # with our inherited permissions system, we decided to rewrite this part to
        # optimize the output.
        # ldif = modlist.modifyModlist(current_entry[0], attr_dict, ignore_oldexistent=1)
        ldif = modifyModlist_finegrained(current_entry[0], attr_dict)

        if not ldif:
            logger.debug("Nothing to update in LDAP")
            return True

        try:
            if new_rdn:
                self.con.rename_s(dn, new_rdn)
                new_base = dn.split(",", 1)[1]
                dn = new_rdn + "," + new_base

            # mod_op : 0 ADD, 1 DELETE, 2 REPLACE
            def encode_values(values: list[str] | str) -> list[bytes]:
                if isinstance(values, list):
                    return [value.encode("utf-8") for value in values]
                elif isinstance(values, str):
                    return [values.encode("utf-8")]
                return values

            encoded_ldif = [
                (mod_op, attribute, encode_values(values))
                for mod_op, attribute, values in ldif
            ]

            self.con.modify_ext_s(dn, encoded_ldif)
        except Exception as e:
            raise MoulinetteError(
                "Error during LDAP update operation:\n"
                f" rdn: {rdn}\n"
                f" attr_dict: {attr_dict}\n"
                f" new_rdn: {new_rdn}\n"
                f" ldif: {ldif}\n"
                f" exception: {e}",
                raw_msg=True,
            )
        else:
            return True

    def validate_uniqueness(self, value_dict: dict[str, str]) -> bool:
        """
        Check uniqueness of values

        Keyword arguments:
            value_dict -- Dictionnary of attributes/values to check

        Returns:
            Boolean | MoulinetteError

        """
        attr_found = self.get_conflict(value_dict)
        if attr_found:
            logger.info(
                "attribute '%s' with value '%s' is not unique",
                attr_found[0],
                attr_found[1],
            )
            raise YunohostError(
                "ldap_attribute_already_exists",
                attribute=attr_found[0],
                value=attr_found[1],
            )
        return True

    def get_conflict(
        self, value_dict: dict[str, str], base_dn: str | None = None
    ) -> tuple[str, str] | None:
        """
        Check uniqueness of values

        Keyword arguments:
            value_dict -- Dictionnary of attributes/values to check

        Returns:
            None | tuple with Fist conflict attribute name and value

        """
        for attr, value in value_dict.items():
            if not self.search(base=base_dn, filter=attr + "=" + value):
                continue
            else:
                return (attr, value)
        return None
