#!/usr/bin/env python3

# Libervia: an XMPP client
# 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 lxml import etree

NS_PIE = "urn:xmpp:pie:0"
NS_ROSTER = "jabber:iq:roster"
NS_PUBSUB = "http://jabber.org/protocol/pubsub"


class PieError(Exception):
    """Error while building or parsing a pie (XEP-0227) document"""


## pie document ##


def build_pie(
    host_jid: str | None, user_name: str, sections: list[etree._Element]
) -> bytes:
    """Build a pie document for a user

    @param host_jid: virtual host JID; the XEP makes the attribute required, so an empty
        value is written when it is unknown (profile JID not set yet)
    @param user_name: node part of the user JID (``<user name=...>``)
    @param sections: section elements to include under the ``<user>`` element
    @return: serialized pie document
    """
    root = etree.Element(f"{{{NS_PIE}}}server-data", nsmap={None: NS_PIE})
    host = etree.SubElement(root, f"{{{NS_PIE}}}host", jid=host_jid or "")
    user = etree.SubElement(host, f"{{{NS_PIE}}}user", name=user_name)
    for section in sections:
        user.append(section)
    return etree.tostring(root, xml_declaration=True, encoding="utf-8", pretty_print=True)


def parse_pie(data: bytes) -> tuple[str, str, list[etree._Element]]:
    """Parse a pie document

    Unknown sections are kept as-is: the XEP requires importers to ignore data they don't
    understand.

    @param data: serialized pie document
    @return: (host jid, user name, sections)
    @raise PieError: if the document is not a valid pie document
    """
    try:
        root = etree.fromstring(data)
    except etree.XMLSyntaxError as e:
        raise PieError(f"invalid pie document: {e}") from e
    if root.tag != f"{{{NS_PIE}}}server-data":
        raise PieError(f"unexpected root element {root.tag!r}, expected pie")
    host = root.find(f"{{{NS_PIE}}}host")
    if host is None:
        raise PieError("missing <host> element in pie document")
    user = host.find(f"{{{NS_PIE}}}user")
    if user is None:
        raise PieError("missing <user> element in pie document")
    return host.get("jid", ""), user.get("name", ""), list(user)


def sections_by_ns(sections: list[etree._Element]) -> dict[str, list[etree._Element]]:
    """Group pie sections by their namespace"""
    by_ns: dict[str, list[etree._Element]] = {}
    for section in sections:
        ns = etree.QName(section).namespace or ""
        by_ns.setdefault(ns, []).append(section)
    return by_ns


## sections: roster ##


def build_roster(
    items: list[etree._Element], version: str | None = None
) -> etree._Element:
    """Build the roster section

    @param items: roster item elements, as in the roster query
    @param version: roster version, if known
    @return: the roster section element
    """
    section = etree.Element(f"{{{NS_ROSTER}}}query", nsmap={None: NS_ROSTER})
    if version:
        section.set("ver", version)
    for item in items:
        section.append(item)
    return section


## sections: pep items ##


def append_xml_payload(parent: etree._Element, data: str | None) -> None:
    """Append item data (stored as XML text) to an element, with text fallback

    @param parent: element to append the payload to
    @param data: payload data (XML string or plain text)
    """
    if not data:
        return
    try:
        parent.append(etree.fromstring(data.encode()))
    except etree.XMLSyntaxError:
        parent.text = data


def item_payload(item: etree._Element) -> str:
    """Serialize the payload of a pubsub item (without the <item> wrapper)

    @param item: item element
    @return: payload as string
    """
    children = list(item)
    if not children:
        return item.text or ""
    # keep the payload in a temporary container (several elements supported)
    tmp = etree.Element("payload")
    if (item.text or "").strip():
        tmp.text = item.text
    for child in children:
        if child.tail and not child.tail.strip():
            child.tail = None
        tmp.append(child)
    # exclusive: only the namespaces actually used by the payload are kept
    data = etree.tostring(tmp, encoding="unicode", exclusive=True)
    return data[len("<payload>") : -len("</payload>")]


def build_pep_items(
    nodes: dict[tuple[str, str | None], list[tuple[str, str, str | None]]],
) -> etree._Element | None:
    """Build the PEP items section

    @param nodes: (node name, service) mapped to the items as (item id, payload,
        publisher) lists, payloads are XML strings or plain text. The service (node owner
        JID) is stored in a ``service`` attribute
    @return: the PEP items section element, or None if there is no node
    """
    if not nodes:
        return None
    section = etree.Element(f"{{{NS_PUBSUB}}}pubsub", nsmap={None: NS_PUBSUB})
    for (node_name, service), items in sorted(
        nodes.items(), key=lambda kv: (kv[0][0], kv[0][1] or "")
    ):
        items_elt = etree.SubElement(section, f"{{{NS_PUBSUB}}}items", node=node_name)
        if service:
            items_elt.set("service", service)
        for item_id, payload, publisher in items:
            item = etree.SubElement(items_elt, f"{{{NS_PUBSUB}}}item", id=item_id)
            if publisher:
                item.set("publisher", publisher)
            append_xml_payload(item, payload)
    return section