"""A lean WS-Management client for Intel AMT.

AMT speaks WS-Management (DMTF DSP0226) with the CIM binding: SOAP 1.2 envelopes POSTed to
``/wsman``, on port 16992 over plain HTTP or 16993 over TLS, authenticated with HTTP Digest.
This client covers the four operations a driver needs: ``Get`` one instance, ``Enumerate``
plus ``Pull`` all instances of a class, and invoke a method. Instances come back as plain
dicts keyed by local element names; a repeated element becomes a list, an element with
children a nested dict, and every value a string, as on the wire.

The envelopes follow Intel's own go-wsman-messages, which every AMT since version 6 accepts.
"""

from __future__ import annotations

import itertools
import xml.etree.ElementTree as ET
from collections.abc import Callable, Iterable, Mapping
from typing import Any
from xml.sax.saxutils import escape, quoteattr

import requests
from requests.auth import HTTPDigestAuth

SOAP_NS = "http://www.w3.org/2003/05/soap-envelope"
ADDRESSING_NS = "http://schemas.xmlsoap.org/ws/2004/08/addressing"
WSMAN_NS = "http://schemas.dmtf.org/wbem/wsman/1/wsman.xsd"
ENUMERATION_NS = "http://schemas.xmlsoap.org/ws/2004/09/enumeration"
TRANSFER_NS = "http://schemas.xmlsoap.org/ws/2004/09/transfer"
ANONYMOUS = ADDRESSING_NS + "/role/anonymous"

PORT = 16992
TLS_PORT = 16993

#: Class name prefix -> the schema its resource URI lives under.
_SCHEMAS = {
    "AMT_": "http://intel.com/wbem/wscim/1/amt-schema/1/",
    "IPS_": "http://intel.com/wbem/wscim/1/ips-schema/1/",
    "CIM_": "http://schemas.dmtf.org/wbem/wscim/1/cim-schema/2/",
}
#: Pages an enumeration may take. AMT's longest lists are a few dozen instances, and it
#: returns all of them in one page when asked for 999; this only stops a runaway loop.
_MAX_PAGES = 100
_HEADERS = {"Content-Type": "application/soap+xml;charset=UTF-8"}

Instance = dict[str, Any]


class WsmanError(Exception):
    """AMT did not answer, or answered with something other than the expected response."""


class WsmanAuthError(WsmanError):
    """AMT refused the credentials (HTTP 401 after the Digest exchange)."""


class WsmanFault(WsmanError):
    """AMT answered with a SOAP fault."""

    def __init__(self, code: str, subcode: str, reason: str, detail: str) -> None:
        self.code = code
        self.subcode = subcode
        self.reason = reason
        self.detail = detail
        super().__init__(f"{subcode or code}: {reason}" + (f" ({detail})" if detail else ""))


def resource_uri(resource: str) -> str:
    """The resource URI of a class, from its prefix (``AMT_``, ``IPS_``, ``CIM_``)."""
    for prefix, base in _SCHEMAS.items():
        if resource.startswith(prefix):
            return base + resource
    raise ValueError(f"{resource!r} is not an AMT, IPS or CIM class")


def _local(tag: str) -> str:
    return tag.rsplit("}", 1)[-1]


def _value(element: ET.Element) -> str | Instance:
    """An element as a string (no children) or a dict of its children."""
    children = list(element)
    if not children:
        return (element.text or "").strip()
    result: Instance = {}
    for child in children:
        key, value = _local(child.tag), _value(child)
        if key not in result:
            result[key] = value
        elif isinstance(result[key], list):
            result[key].append(value)
        else:
            result[key] = [result[key], value]
    return result


def to_instance(element: ET.Element | None) -> Instance:
    """An instance element as a dict; an empty or missing one as an empty dict."""
    if element is None:
        return {}
    value = _value(element)
    return value if isinstance(value, dict) else {}


def _fault(element: ET.Element) -> WsmanFault:
    def text(path: str) -> str:
        found = element.find(path)
        return "".join(found.itertext()).strip() if found is not None else ""

    def unprefixed(qname: str) -> str:
        return qname.split(":", 1)[-1]

    s = f"{{{SOAP_NS}}}"
    return WsmanFault(
        code=unprefixed(text(f"{s}Code/{s}Value")),
        subcode=unprefixed(text(f"{s}Code/{s}Subcode/{s}Value")),
        reason=text(f"{s}Reason/{s}Text"),
        detail=text(f"{s}Detail"),
    )


class WsmanTransport:
    """WS-Management over HTTP(S) with Digest auth, to one AMT.

    ``tls`` defaults to on for port 16993 and off otherwise. ``verify`` is passed to
    ``requests``: ``True`` checks the certificate against its CA bundle, a path against
    that CA file, ``False`` not at all. ``retries`` counts extra attempts after a
    connection error; a response, even an error, is never retried. ``recorder``, when set,
    is called with the class, the operation and the raw body of every answer (to record
    fixtures).
    """

    def __init__(
        self,
        host: str,
        username: str,
        password: str,
        port: int | None = None,
        tls: bool | None = None,
        verify: bool | str = True,
        timeout: float = 60,
        retries: int = 1,
    ) -> None:
        self.port = int(port or PORT)
        self.tls = self.port == TLS_PORT if tls is None else bool(tls)
        self.url = f"{'https' if self.tls else 'http'}://{host}:{self.port}/wsman"
        self.timeout = timeout
        self.retries = max(0, int(retries))
        self.session = requests.Session()
        self.session.auth = HTTPDigestAuth(username, password)
        self.session.verify = verify
        self._message_ids = itertools.count(1)
        self.recorder: Callable[[str, str, bytes], None] | None = None

    def close(self) -> None:
        self.session.close()

    # --- operations ----------------------------------------------------------------------

    def get(self, resource: str, selectors: Mapping[str, str] | None = None) -> Instance:
        """The one instance of ``resource`` (or the one ``selectors`` pick)."""
        body = self._post(resource, f"{TRANSFER_NS}/Get", "", selectors)
        return to_instance(next(iter(body), None))

    def enumerate(self, resource: str) -> list[Instance]:
        """Every instance of ``resource``: ``Enumerate``, then ``Pull`` to the end."""
        e = f"{{{ENUMERATION_NS}}}"
        body = self._post(
            resource, f"{ENUMERATION_NS}/Enumerate", f'<Enumerate xmlns="{ENUMERATION_NS}"/>'
        )
        context = body.findtext(f"{e}EnumerateResponse/{e}EnumerationContext")
        instances: list[Instance] = []
        for _ in range(_MAX_PAGES):
            if not context:
                return instances
            body = self._post(
                resource,
                f"{ENUMERATION_NS}/Pull",
                f'<Pull xmlns="{ENUMERATION_NS}"><EnumerationContext>{escape(context)}'
                f"</EnumerationContext><MaxElements>999</MaxElements>"
                f"<MaxCharacters>99999</MaxCharacters></Pull>",
            )
            page = body.find(f"{e}PullResponse")
            if page is None:
                raise WsmanError(f"{resource}: Pull answered without a PullResponse")
            instances.extend(to_instance(item) for item in page.findall(f"{e}Items/*"))
            if page.find(f"{e}EndOfSequence") is not None:
                return instances
            context = page.findtext(f"{e}EnumerationContext")
        raise WsmanError(f"{resource}: no end of the enumeration after {_MAX_PAGES} pages")

    def invoke(
        self,
        resource: str,
        method: str,
        params: Iterable[tuple[str, Any]] = (),
        raw: str = "",
        selectors: Mapping[str, str] | None = None,
        check: bool = True,
    ) -> Instance:
        """Call ``method`` of ``resource`` and return its output parameters.

        ``params`` are sent in order as ``<h:Name>value</h:Name>``; ``raw`` is appended
        unescaped after them, for parameters that are references (``h:`` is the class's
        namespace). With ``check``, a ``ReturnValue`` other than 0 raises.
        """
        uri = resource_uri(resource)
        inner = "".join(f"<h:{name}>{escape(str(value))}</h:{name}>" for name, value in params)
        body = self._post(
            resource,
            f"{uri}/{method}",
            f'<h:{method}_INPUT xmlns:h="{uri}">{inner}{raw}</h:{method}_INPUT>',
            selectors,
        )
        output = to_instance(body.find(f"{{{uri}}}{method}_OUTPUT"))
        code = output.get("ReturnValue")
        if check and code not in (None, "0"):
            raise WsmanError(f"{resource}.{method} failed: ReturnValue {code}")
        return output

    # --- plumbing ------------------------------------------------------------------------

    def _envelope(
        self, uri: str, action: str, body: str, selectors: Mapping[str, str] | None
    ) -> str:
        selector_set = ""
        if selectors:
            selector_set = (
                "<w:SelectorSet>"
                + "".join(
                    f"<w:Selector Name={quoteattr(name)}>{escape(value)}</w:Selector>"
                    for name, value in selectors.items()
                )
                + "</w:SelectorSet>"
            )
        return (
            f'<?xml version="1.0" encoding="utf-8"?>'
            f'<Envelope xmlns="{SOAP_NS}" xmlns:a="{ADDRESSING_NS}" xmlns:w="{WSMAN_NS}">'
            f"<Header><a:Action>{action}</a:Action><a:To>/wsman</a:To>"
            f"<w:ResourceURI>{uri}</w:ResourceURI>"
            f"<a:MessageID>{next(self._message_ids)}</a:MessageID>"
            f"<a:ReplyTo><a:Address>{ANONYMOUS}</a:Address></a:ReplyTo>"
            f"<w:OperationTimeout>PT60S</w:OperationTimeout>{selector_set}</Header>"
            f"<Body>{body}</Body></Envelope>"
        )

    def _post(
        self, resource: str, action: str, body: str, selectors: Mapping[str, str] | None = None
    ) -> ET.Element:
        envelope = self._envelope(resource_uri(resource), action, body, selectors)
        response = self._send(envelope.encode())
        if self.recorder is not None:
            self.recorder(resource, action.rsplit("/", 1)[-1], response.content)
        return self._body(response)

    def _send(self, data: bytes) -> requests.Response:
        error: requests.RequestException | None = None
        for _ in range(self.retries + 1):
            try:
                return self.session.post(
                    self.url, data=data, headers=_HEADERS, timeout=self.timeout
                )
            except requests.RequestException as exc:
                error = exc
        raise WsmanError(f"{self.url}: {error}") from error

    def _body(self, response: requests.Response) -> ET.Element:
        """The SOAP body of a response; a fault, a refusal or anything else raises."""
        if response.status_code == 401:
            raise WsmanAuthError(f"{self.url}: credentials refused (HTTP 401)")
        body = None
        # A DTD has no business in a WS-Management response; refusing one keeps entity
        # expansion out of the parser.
        if b"<!DOCTYPE" not in response.content:
            try:
                body = ET.fromstring(response.content).find(f"{{{SOAP_NS}}}Body")
            except ET.ParseError:
                body = None
        if body is not None:
            fault = body.find(f"{{{SOAP_NS}}}Fault")
            if fault is not None:
                raise _fault(fault)
        if body is None or response.status_code >= 400:
            raise WsmanError(f"{self.url}: HTTP {response.status_code}, no WS-Management answer")
        return body