#!/usr/bin/env python3

# Libervia plugin for Pubsub Contacts.
# Copyright (C) 2009-2026 Jérôme Poisson (goffi@goffi.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/>.

from dataclasses import dataclass, field
from typing import Any, Iterable, Iterator, TYPE_CHECKING, Self, cast

from pydantic import BaseModel, Field, RootModel, field_validator
import shortuuid
from twisted.internet import defer
from twisted.words.protocols.jabber.jid import JID
from twisted.words.protocols.jabber.xmlstream import XMPPHandler
from twisted.words.xish import domish
from tx_xmpp import disco
from tx_xmpp.pubsub import NS_PUBSUB, NS_PUBSUB_EVENT
from zope.interface import implementer

from libervia.backend.core import exceptions
from libervia.backend.core.constants import Const as C
from libervia.backend.core.core_types import SatXMPPEntity
from libervia.backend.core.i18n import _
from libervia.backend.core.log import getLogger
from libervia.backend.plugins.plugin_xep_0059 import XEP_0059, RSMRequest, RSMResponse
from libervia.backend.plugins.plugin_xep_0060 import XEP_0060
from libervia.backend.plugins.plugin_xep_0163 import XEP_0163
from libervia.backend.tools.xml_tools import find_ancestor

if TYPE_CHECKING:
    from libervia.backend.core.main import LiberviaBackend
    from tx_xmpp.pubsub import ItemsEvent

log = getLogger(__name__)


PLUGIN_INFO = {
    C.PI_NAME: "Pubsub Contacts",
    C.PI_IMPORT_NAME: "CONTACTS",
    C.PI_TYPE: "XEP",
    C.PI_MODES: C.PLUG_MODE_BOTH,
    C.PI_PROTOCOLS: [],
    C.PI_DEPENDENCIES: ["XEP-0059", "XEP-0060", "XEP-0163"],
    C.PI_RECOMMENDATIONS: [],
    C.PI_MAIN: "PubsubContacts",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _("""Handle Contacts Metadata on Pubsub."""),
}

NS_CONTACTS_BASE = "urn:xmpp:contacts"
NS_CONTACTS = f"{NS_CONTACTS_BASE}:0"
NS_CONTACTS_GROUPS = "urn:xmpp:contacts-groups"


class Group(BaseModel):

    id: str | None = None
    name: str | None = None
    description: str | None = None
    encrypted: bool = True

    @classmethod
    def from_element(cls, element: domish.Element, metadata: dict | None = None) -> Self:
        """Parse from <group> element.

        @param element: <group> element or a direct parent.
        """
        if element.uri != NS_CONTACTS or element.name != "group":
            try:
                element = next(element.elements(NS_CONTACTS, "group"))
            except StopIteration:
                raise ValueError(f"Invalid element: {element.toXml()}")

        group_id = element.getAttribute("id")
        if not group_id:
            raise ValueError(f"Invalid element, missing id: {element.toXml()}")

        encrypted = False
        if metadata is not None:
            encrypted = metadata.get("encrypted", {})
            if group_id in encrypted:
                encrypted = True

        name = element.getAttribute("name")

        for child in element.elements(NS_CONTACTS, "description"):
            description = str(child)
            break
        else:
            description = None

        return cls(id=group_id, name=name, description=description, encrypted=encrypted)

    def to_element(self, id_only: bool = False) -> domish.Element:
        """Convert to <group> element.

        @param id_only: if True, only use the "id" attribute (used for <contact/> element)
        """
        group_elt = domish.Element((NS_CONTACTS, "group"))
        if not self.id:
            raise ValueError("self.id must be set.")
        group_elt["id"] = self.id
        if not id_only:
            if not self.name:
                raise ValueError("self.name must be set.")
            group_elt["name"] = self.name

            if self.description is not None:
                group_elt.addElement("description", content=self.description)

        return group_elt


class Groups(RootModel):
    root: list[Group]

    def __iter__(self) -> Iterator[Group]:  # type: ignore
        return iter(self.root)

    def __getitem__(self, item) -> Group:
        return self.root[item]

    def __len__(self) -> int:
        return len(self.root)

    def append(self, item: Group) -> None:
        self.root.append(item)

    def extend(self, items: Iterable[Group]) -> None:
        self.root.extend(items)

    def remove(self, item: Group) -> None:
        self.root.remove(item)

    def clear(self) -> None:
        self.root.clear()

    def sort(self, key=None, reverse=False) -> None:
        self.root.sort(key=key, reverse=reverse)  # type: ignore


class Contact(BaseModel):

    id: str | None = None
    name: str | None = None
    identities: dict[str, str]
    groups: Groups = Field(default_factory=lambda: Groups([]))
    encrypted: bool = True

    @field_validator("identities", mode="after")
    @classmethod
    def identities_is_not_empty(cls, value: dict[str, str]) -> dict[str, str]:
        if not value:
            raise ValueError('"identities" field must have at least one value.')
        if len(value) > 1 or (next(iter(value)) != "jid"):
            raise NotImplementedError(f'Only "jid" type is supported for now.')
        return value

    @classmethod
    def from_element(
        cls,
        element: domish.Element,
        contacts_data: "ContactsData",
        metadata: dict | None = None,
    ) -> "Contact":
        """Parse from <contact> element.

        @param element: <contact> element or a direct parent.
        @param contacts_data: used to fill group data, as in <contact> we have only the
            group id.
        @param metadata: metadata as returned by "XEP_0060.get_items".
        """
        if element.uri != NS_CONTACTS or element.name != "contact":
            try:
                element = next(element.elements(NS_CONTACTS, "contact"))
            except StopIteration:
                raise ValueError(f"Invalid element: {element.toXml()}")

        data: dict[str, Any] = {"encrypted": False}

        name = element.getAttribute("name")
        if name is not None:
            data["name"] = name

        identities = {}

        for child in element.elements():
            if child.uri != NS_CONTACTS:
                continue

            match child.name:
                case "identity":
                    identity_type = child.getAttribute("type")
                    if not identity_type:
                        raise ValueError(
                            "<identity> element must have a 'type' attribute"
                        )
                    identities[identity_type] = str(child)

                case "group":
                    group = Group.from_element(child)
                    if not group.id:
                        log.warning(f'Invalid group, no "id", ignoring: {child.toXml()}')
                        continue
                    if "groups" not in data:
                        data["groups"] = Groups(root=[])
                    try:
                        existing_group = contacts_data.ids_to_groups[group.id]
                    except KeyError:
                        log.warning(f"Unknown group, ignoring: {group}")
                        continue
                    else:
                        group = existing_group
                    data["groups"].append(group)

        if not identities:
            raise exceptions.DataError(
                "<contact> element must have at least one <identity> element."
            )

        data["identities"] = identities
        try:
            item_elt = find_ancestor(element, "item", [NS_PUBSUB, NS_PUBSUB_EVENT])
        except exceptions.NotFound:
            log.warning("Could not find <item> ancestor of <contact>")
        else:
            item_id = item_elt["id"]
            data["id"] = item_id
            if metadata is not None:
                encrypted = metadata.get("encrypted", {})
                if item_id in encrypted:
                    data["encrypted"] = True

        return cls(**data)

    def to_element(self) -> domish.Element:
        """Convert to <contact> element."""
        contact_elt = domish.Element((NS_CONTACTS, "contact"))

        if self.name is not None:
            contact_elt["name"] = self.name

        for identity_type, identity_value in self.identities.items():
            identity_elt = contact_elt.addElement("identity", content=identity_value)
            identity_elt["type"] = identity_type

        for group in self.groups:
            contact_elt.addChild(group.to_element(id_only=True))

        return contact_elt


class Contacts(RootModel):
    root: list[Contact]

    def __iter__(self) -> Iterator[Contact]:  # type: ignore
        return iter(self.root)

    def __getitem__(self, item) -> Contact:
        return self.root[item]

    def __len__(self) -> int:
        return len(self.root)

    def append(self, item: Contact) -> None:
        self.root.append(item)

    def extend(self, items: Iterable[Contact]) -> None:
        self.root.extend(items)

    def remove(self, item: Contact) -> None:
        self.root.remove(item)

    def clear(self) -> None:
        self.root.clear()

    def sort(self, key=None, reverse=False) -> None:
        self.root.sort(key=key, reverse=reverse)  # type: ignore


class ContactsRequest(BaseModel):
    clear: bool = Field(
        default=True, description="If true, retrieve clear text contacts."
    )
    encrypted: bool = Field(
        default=True, description="If true, retrieve end-to-end encrypted text contacts."
    )
    ignore_cache: bool = Field(
        default=False,
        description=(
            "If True, won't use cached value and will fetch from pubsub service again."
        ),
    )


class ContactSetRequest(BaseModel):
    contact: Contact


@dataclass
class ContactsData:
    contacts: Contacts = field(default_factory=lambda: Contacts([]))
    groups: Groups = field(default_factory=lambda: Groups([]))
    _jid_to_contacts: dict[str, Contact] | None = None
    _ids_to_groups: dict[str, Group] | None = None

    @property
    def jid_to_contacts(self) -> dict[str, Contact]:
        """A map from contact JID to Contact instance.

        A str is used instead of jid.JID for convenience, as for now Contact.identities
        only support str to str mapping.
        """
        if self._jid_to_contacts is None:
            if self.contacts is None:
                return {}
            jid_to_contacts = {}
            for contact in self.contacts:
                if not contact.id:
                    log.warning(f"ignoring contact with no id: {contact}")
                    continue
                jid = contact.identities.get("jid")
                if not jid:
                    log.warning(f"No JID found for {contact}, ignoring")
                    continue
                jid_to_contacts[jid] = contact
            self._jid_to_contacts = jid_to_contacts
        return self._jid_to_contacts

    @property
    def ids_to_groups(self) -> dict[str, Group]:
        """A map from group IDs to Group instance."""
        if self._ids_to_groups is None:
            if self.groups is None:
                return {}
            ids_to_groups = {}
            for group in self.groups:
                if group.id is None:
                    log.warning(f'"id" not set for this group, ignoring: {group}')
                    continue
                ids_to_groups[group.id] = group
            self._ids_to_groups = ids_to_groups
        return self._ids_to_groups

    def reset(self) -> None:
        """Reset data internal values.

        To be used when groups and contacts are fetched again.
        """
        self._jid_to_contacts = self._ids_to_groups = None


class PubsubContacts:
    namespace = NS_CONTACTS

    def __init__(self, host: "LiberviaBackend") -> None:
        log.info(f"plugin {PLUGIN_INFO[C.PI_NAME]!r} initialization")
        self.host = host
        self._p = cast(XEP_0060, host.plugins["XEP-0060"])
        self._rsm = cast(XEP_0059, host.plugins["XEP-0059"])
        pep = cast(XEP_0163, host.plugins["XEP-0163"])
        pep.add_pep_event(None, NS_CONTACTS_BASE, self._on_contacts_update)
        pep.add_pep_event(None, NS_CONTACTS_GROUPS, self._on_groups_update)
        host.bridge.add_method(
            "ps_contacts_get",
            ".plugin",
            in_sign="ss",
            out_sign="s",
            method=self._contacts_get,
            async_=True,
        )
        host.bridge.add_method(
            "ps_contact_set",
            ".plugin",
            in_sign="ss",
            out_sign="s",
            method=self._contact_set,
            async_=True,
        )

    def get_handler(self, client: SatXMPPEntity) -> "PubsubContactsHandler":
        return PubsubContactsHandler()

    def _on_contacts_update(self, itemsEvent, profile: str) -> None:
        client = self.host.get_client(profile)
        data = client.get_plugin_data(C.PI_IMPORT_NAME, ContactsData)
        for item_elt in itemsEvent.items:
            if item_elt.name == "retract":
                item_id = item_elt["id"]
                try:
                    to_delete = next(c for c in data.contacts if c.id == item_id)
                except StopIteration:
                    pass
                else:
                    data.contacts.remove(to_delete)
                    data.reset()
            else:
                item_id = item_elt["id"]
                if item_elt.hasAttribute("encrypted"):
                    metadata = {"encrypted": {item_id: item_elt.encrypted}}
                else:
                    metadata = None

                contact = Contact.from_element(item_elt, data, metadata)
                try:
                    existing_contact = next(c for c in data.contacts if c.id == item_id)
                except StopIteration:
                    data.contacts.append(contact)
                    data.reset()
                else:
                    # We only replace the contact if they differs
                    if contact != existing_contact:
                        data.contacts.remove(existing_contact)
                        data.contacts.append(contact)
                        data.reset()

    def _on_groups_update(self, itemsEvent, profile: str) -> None:
        client = self.host.get_client(profile)
        data = client.get_plugin_data(C.PI_IMPORT_NAME, ContactsData)
        for item_elt in itemsEvent.items:
            if item_elt.name == "retract":
                item_id = item_elt["id"]
                try:
                    to_delete = next(g for g in data.groups if g.id == item_id)
                except StopIteration:
                    pass
                else:
                    data.groups.remove(to_delete)
                    data.reset()
            else:
                if item_elt.hasAttribute("encrypted"):
                    metadata = {"encrypted": {item_elt["id"]: item_elt.encrypted}}
                else:
                    metadata = None

                group = Group.from_element(item_elt, metadata)
                assert group.id is not None
                try:
                    existing_group = data.ids_to_groups[group.id]
                except KeyError:
                    data.groups.append(group)
                    data.reset()
                else:
                    # We only replace the group if they differs
                    if group != existing_group:
                        data.groups.remove(existing_group)
                        data.groups.append(group)
                        data.reset()

    def _contacts_get(
        self, contacts_request_s: str, profile_key: str
    ) -> defer.Deferred[str]:
        client = self.host.get_client(profile_key)
        contacts_request = ContactsRequest.model_validate_json(contacts_request_s)
        d = defer.ensureDeferred(self.contacts_get(client, contacts_request))
        d.addCallback(lambda contacts: contacts.model_dump_json())
        d = cast(defer.Deferred[str], d)
        return d

    async def contacts_get(
        self, client: SatXMPPEntity, contacts_request: ContactsRequest
    ) -> Contacts:
        """Retrieve contacts from pubsub.

        This method fetches contacts and groups from the pubsub service.
        If contacts_request.ignore_cache is not set, cached version will be used if it
        exists.
        @param client: client instance
        @param contacts_request: Request parameters
        @return: Retrieved contacts
        """
        data = client.get_plugin_data(C.PI_IMPORT_NAME, ContactsData)
        if contacts_request.ignore_cache:
            data.contacts.clear()
            data.groups.clear()
            data.reset()
        else:
            if data.contacts:
                return data.contacts
        # We first retrieve groups
        rsm_request = RSMRequest(max=20)
        while True:
            try:
                groups_items, groups_metadata = await self._p.get_items(
                    client, None, NS_CONTACTS_GROUPS, rsm_request=rsm_request
                )
            except exceptions.NotFound:
                log.debug("Group node doesn't exist.")
                break
            except Exception as e:
                log.warning(f"Can't retrieve groups: {e}")
                break
            else:
                data.groups.extend(
                    Group.from_element(e, groups_metadata) for e in groups_items
                )
                if rsm_dict := groups_metadata.get("rsm"):
                    rsm_response = RSMResponse(**rsm_dict)
                    rsm_request = self._rsm.get_next_request(
                        rsm_request, rsm_response, False
                    )
                    if rsm_request is None:
                        break
                else:
                    break

        # Then contacts
        rsm_request = RSMRequest(max=1)
        while True:
            try:
                contacts_items, metadata = await self._p.get_items(
                    client, None, NS_CONTACTS_BASE, rsm_request=rsm_request
                )
            except exceptions.NotFound:
                log.debug("Contacts node doesn't exist.")
                break
            except Exception as e:
                log.warning(f"Can't retrieve contacts: {e}")
                break
            else:
                try:
                    new_contacts = Contacts(
                        [Contact.from_element(e, data, metadata) for e in contacts_items]
                    )
                except Exception as e:
                    log.warning(f"Can't parse contacts: {e}")
                    raise e
                data.contacts.extend(new_contacts)
                if rsm_dict := metadata.get("rsm"):
                    rsm_response = RSMResponse(**rsm_dict)
                    rsm_request = self._rsm.get_next_request(
                        rsm_request, rsm_response, False
                    )
                    if rsm_request is None:
                        break
                else:
                    break
        return data.contacts

    def _contact_set(
        self, contacts_set_request_s: str, profile_key: str
    ) -> defer.Deferred[str]:
        client = self.host.get_client(profile_key)
        contacts_set_request = ContactSetRequest.model_validate_json(
            contacts_set_request_s
        )
        d = defer.ensureDeferred(self.contacts_set(client, contacts_set_request))
        return d

    async def contacts_set(
        self, client: SatXMPPEntity, contacts_set_request: ContactSetRequest
    ) -> str:
        """Add or update contact.

        This method will add the contact to the contacts node, creating it if necessary.
        It will first check if the contact already exist, and update the relevant item in
        this case.
        Similarly, contact groups will be created or used if they already exist, and the
        groups node will be created if necessary.

        @param client: client instance
        @param contacts_set_request: Set request data.
        @return: contact item ID
        """
        # we need to fetch contacts to match existing ones (or groups).
        await self.contacts_get(client, ContactsRequest())
        data = client.get_plugin_data(C.PI_IMPORT_NAME, ContactsData)
        contact = contacts_set_request.contact

        if contact.groups:
            new_groups = Groups([])
            for group in contact.groups:
                if group.id:
                    try:
                        existing_group = data.ids_to_groups[group.id]
                    except KeyError:
                        log.warning(f"Can't find group with id {group.id}, ignoring.")
                        continue
                elif group.name:
                    group.name = group.name.strip()
                    if not group.name:
                        log.warning("A group name can't only contain spaces, ignoring.")
                        continue
                    try:
                        existing_group = next(
                            g
                            for g in data.groups
                            if g.name and g.name.lower() == group.name.lower()
                        )
                    except StopIteration:
                        log.debug(f"Adding new group: {group.name!r}")
                        new_id = shortuuid.uuid()
                        group.id = new_id
                        extra = {
                            self._p.EXTRA_AUTOCREATE: True,
                            self._p.EXTRA_PUBLISH_OPTIONS: {
                                XEP_0060.OPT_ACCESS_MODEL: self._p.ACCESS_WHITELIST
                            },
                        }
                        if contact.encrypted:
                            extra["encrypted"] = True
                        item_id = await self._p.send_item(
                            client,
                            None,
                            NS_CONTACTS_GROUPS,
                            group.to_element(),
                            item_id=new_id,
                            extra=extra,
                        )
                        existing_group = group
                else:
                    log.warning(f"Group has no name nor id, ignoring: {group!r}.")
                    continue
                new_groups.append(existing_group)
            contact.groups = new_groups

        if contact.id is None:
            try:
                contact_jid = contact.identities["jid"]
            except KeyError:
                raise NotImplementedError("Contact has no JID, this is not supported yet")
            existing_contact = data.jid_to_contacts.get(contact_jid)
            if existing_contact is not None:
                contact.id = existing_contact.id
        contact_elt = contact.to_element()
        extra = {
            self._p.EXTRA_AUTOCREATE: True,
            self._p.EXTRA_PUBLISH_OPTIONS: {
                XEP_0060.OPT_ACCESS_MODEL: self._p.ACCESS_WHITELIST
            },
        }
        if contact.encrypted:
            extra["encrypted"] = True
        item_id = await self._p.send_item(
            client,
            None,
            NS_CONTACTS_BASE,
            contact_elt,
            item_id=contact.id,
            extra=extra,
        )
        assert item_id is not None
        return item_id


@implementer(disco.IDisco)
class PubsubContactsHandler(XMPPHandler):

    def getDiscoInfo(
        self, requestor: JID, target: JID, nodeIdentifier: str = ""
    ) -> list[disco.DiscoFeature]:
        """Get disco info for pubsub extended discovery

        @param requestor: JID of the requesting entity
        @param target: JID of the target entity
        @param nodeIdentifier: optional node identifier
        @return: list of disco features
        """
        return [
            disco.DiscoFeature(NS_CONTACTS),
        ]

    def getDiscoItems(
        self, requestor: JID, target: JID, nodeIdentifier: str = ""
    ) -> list[disco.DiscoItem]:
        """Get disco items with extended discovery support

        @param requestor: JID of the requesting entity
        @param target: JID of the target entity
        @param nodeIdentifier: optional node identifier
        @return: list of disco items
        """
        return []