"""Connection helper for the M-Files gRPC API.

It does the parts of the API that the .proto files cannot express:

* the channel: TLS (optionally trusting a given CA certificate) or insecure, with the message size
  limits of M-Files' own client (responses of any size, requests up to 48 MiB);
* call metadata: ``mfiles-is-remote-call``, a fresh ``mfiles-activity-id-bin`` on every call
  and, once logged in, ``mfiles-session-id-bin``;
* cookies: the ``mfilesmsm`` cookie that M-Files Cloud sets goes back on every later call;
* logging in with an M-Files user name and password or with an SSO token, and logging out;
* decoding errors: a failed call raises :class:`MFilesError` with the M-Files error code and
  message, read from the ``mfiles-error-*-bin`` trailers;
* the vault's metadata structure version from the ``mfiles-metadata-version-bin`` trailer.

Configuration comes from environment variables (see :meth:`Config.from_env`).

Usage::

    from mfiles_client import MFilesClient

    with MFilesClient.from_env() as client:   # connects and logs in
        search = client.stub(IRPCSearch_pb2_grpc.IRPCSearchStub)
        ...
"""

from __future__ import annotations

import os
import struct
import sys
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import grpc

# The stubs are generated into gen/ by generate.py.
sys.path.insert(0, str(Path(__file__).resolve().parent / "gen"))

from mfiles import IRPCLogin_pb2, IRPCLogin_pb2_grpc, enums_pb2, structs_pb2  # noqa: E402

# Message size limits, the same as M-Files' own client: accept responses of any size (-1), and
# send at most what the server accepts by default (48 MiB).
MAX_RECEIVE_MESSAGE_BYTES = -1
MAX_SEND_MESSAGE_BYTES = 48 * 1024 * 1024

# Call metadata. gRPC base64-encodes the values of "-bin" keys itself: pass raw bytes.
REMOTE_CALL_KEY = "mfiles-is-remote-call"
ACTIVITY_ID_KEY = "mfiles-activity-id-bin"
SESSION_ID_KEY = "mfiles-session-id-bin"
COOKIE_KEY = "cookie"
SET_COOKIE_KEY = "set-cookie"
# Trailers.
ERROR_STACK_KEY = "mfiles-error-serialized-bin"
ERROR_TEXT_KEY = "mfiles-error-text-bin"
METADATA_VERSION_KEY = "mfiles-metadata-version-bin"

# The generic code of an error that is only passed up the stack; it says nothing of the cause.
E_PASS_THROUGH = 0x80040001


@dataclass
class Config:
    """Where and how to connect. Its repr leaves out the password and the token."""

    address: str                        # host:port
    vault: str                          # vault GUID with braces
    tls: bool = True
    username: str = ""
    password: str = field(default="", repr=False)
    token: str = field(default="", repr=False)  # SSO: token from the vault's identity provider
    auth_plugin: str = ""               # SSO: the authentication plugin's name
    auth_scope: str = ""                # SSO: the plugin's configuration scope
    ca_cert: str = ""                   # PEM file of root certificates to trust

    @classmethod
    def from_env(cls, require_vault: bool = True) -> "Config":
        def env(name: str, default: str = "") -> str:
            return os.environ.get(name, default)

        required = ("MFILES_ADDRESS", "MFILES_VAULT") if require_vault else ("MFILES_ADDRESS",)
        missing = [n for n in required if not env(n)]
        if missing:
            raise SystemExit(f"Set the environment variable(s): {', '.join(missing)}")
        return cls(
            address=env("MFILES_ADDRESS"),
            vault=env("MFILES_VAULT"),
            tls=env("MFILES_TLS", "true").strip().lower() not in ("false", "0", "no"),
            username=env("MFILES_USERNAME"),
            password=env("MFILES_PASSWORD"),
            token=env("MFILES_TOKEN"),
            auth_plugin=env("MFILES_AUTH_PLUGIN"),
            auth_scope=env("MFILES_AUTH_SCOPE"),
            ca_cert=env("MFILES_CA_CERT"),
        )

    @property
    def host(self) -> str:
        return self.address.rsplit(":", 1)[0].strip("[]")


class MFilesError(grpc.RpcError):
    """A failed call, with the M-Files error decoded from the trailers.

    ``hresult`` is the M-Files error code (an HRESULT) as an unsigned 32-bit number: the code of
    the top-level error, the last entry of the error stack, or of the last entry that is not the
    generic 0x80040001. ``hresult_hex`` formats it (``0x8004000B``) and ``hresult_signed`` gives
    it as a signed integer. ``codes`` holds every code in the stack, from the original cause
    (first) to the top-level error (last); ``has(code)`` tests whether a code is anywhere in it.
    ``hresult`` is None when the server sent no error stack, e.g. for UNAUTHENTICATED.
    ``status`` is the gRPC status code and ``message`` the status message.
    """

    def __init__(self, error: grpc.RpcError):
        super().__init__()
        self.rpc_error = error
        self.status: grpc.StatusCode = error.code()
        self.message: str = error.details() or ""
        trailers = dict(error.trailing_metadata() or ())
        # [START decode-error]
        self.long_message: str = trailers.get(ERROR_TEXT_KEY, b"").decode("utf-8", "replace")
        self.stack = list(structs_pb2.RPCErrorInfoArray.FromString(
            trailers.get(ERROR_STACK_KEY, b"")).value)
        self.codes = [entry.error_code & 0xFFFFFFFF for entry in self.stack]  # cause first
        # The top-level error, skipping entries that only pass the error on.
        self.hresult: Optional[int] = next(
            (code for code in reversed(self.codes) if code != E_PASS_THROUGH),
            self.codes[-1] if self.codes else None)
        # [END decode-error]

    def has(self, code: int) -> bool:
        """True if the error stack contains the M-Files error code, e.g. 0x8004000B."""
        return code in self.codes

    @property
    def hresult_signed(self) -> Optional[int]:
        if self.hresult is None:
            return None
        return self.hresult - (1 << 32) if self.hresult >= 1 << 31 else self.hresult

    @property
    def hresult_hex(self) -> str:
        return "" if self.hresult is None else f"0x{self.hresult:08X}"

    def code(self) -> grpc.StatusCode:
        return self.status

    def details(self) -> str:
        return self.message

    def __str__(self) -> str:
        if self.hresult is None:
            return f"gRPC error {self.status.name}: {self.message}"
        return f"M-Files error {self.hresult_hex} ({self.status.name}): {self.message}"


class _MetadataInterceptor(grpc.UnaryUnaryClientInterceptor,
                           grpc.StreamUnaryClientInterceptor):
    """Adds the M-Files call metadata, keeps cookies and trailers, and decodes errors."""

    def __init__(self, client: "MFilesClient"):
        self._client = client

    def _details(self, details: grpc.ClientCallDetails) -> grpc.ClientCallDetails:
        metadata = list(details.metadata or [])
        metadata.append((REMOTE_CALL_KEY, "true"))
        metadata.append((ACTIVITY_ID_KEY, uuid.uuid4().bytes))
        if self._client.session_id:
            metadata.append((SESSION_ID_KEY, self._client.session_id))
        if self._client.cookies:
            metadata.append((COOKIE_KEY, "; ".join(
                f"{name}={value}" for name, value in self._client.cookies.items())))
        return details._replace(metadata=metadata)

    def _complete(self, outcome):
        try:
            outcome.result()
        except grpc.RpcError as error:
            self._client._record_metadata(error.initial_metadata(), error.trailing_metadata())
            raise MFilesError(error) from None
        self._client._record_metadata(outcome.initial_metadata(), outcome.trailing_metadata())
        return outcome

    def intercept_unary_unary(self, continuation, client_call_details, request):
        return self._complete(continuation(self._details(client_call_details), request))

    def intercept_stream_unary(self, continuation, client_call_details, request_iterator):
        # Client-streaming calls (a stream of requests, one response).
        return self._complete(continuation(self._details(client_call_details), request_iterator))


class MFilesClient:
    """One channel to an M-Files server, logged in to one vault."""

    def __init__(self, config: Config, client_name: str = ""):
        self.config = config
        # Shown in the server's session list; defaults to the program's name.
        self.client_name = client_name or Path(sys.argv[0]).stem or "python"
        self.session_id: bytes = b""
        self.cookies: dict[str, str] = {}  # Cookies the server set, such as mfilesmsm.
        self.session: Optional[structs_pb2.SessionData] = None
        self.metadata_structure_version: Optional[int] = None
        # [START channel]
        options = [("grpc.max_receive_message_length", MAX_RECEIVE_MESSAGE_BYTES),
                   ("grpc.max_send_message_length", MAX_SEND_MESSAGE_BYTES)]
        if config.tls:
            root_certificates = Path(config.ca_cert).read_bytes() if config.ca_cert else None
            credentials = grpc.ssl_channel_credentials(root_certificates=root_certificates)
            self._channel = grpc.secure_channel(config.address, credentials, options=options)
        else:
            self._channel = grpc.insecure_channel(config.address, options=options)
        self.channel = grpc.intercept_channel(self._channel, _MetadataInterceptor(self))
        # [END channel]
        self.login = IRPCLogin_pb2_grpc.IRPCLoginStub(self.channel)

    @classmethod
    def from_env(cls, client_name: str = "", log_in: bool = True) -> "MFilesClient":
        """Connects with the configuration in the environment and, by default, logs in.

        Without log_in, MFILES_VAULT is optional: set config.vault before logging in.
        """
        client = cls(Config.from_env(require_vault=log_in), client_name)
        if log_in:
            try:
                client.log_in()
            except BaseException:
                client.close()
                raise
        return client

    def stub(self, stub_class):
        """Returns a stub of a service: ``client.stub(IRPCSearch_pb2_grpc.IRPCSearchStub)``."""
        return stub_class(self.channel)

    def log_in(self) -> structs_pb2.SessionData:
        """Logs in to the vault with a token if one is configured, else with user and password."""
        if self.config.token:
            return self.log_in_with_token(self.config.token, self.config.auth_plugin,
                                          self.config.auth_scope)
        return self.log_in_with_password(self.config.username, self.config.password)

    def log_in_with_password(self, username: str, password: str) -> structs_pb2.SessionData:
        # [START login-password]
        authentication = structs_pb2.AuthDataClient(
            type=enums_pb2.AUTH_DATA_TYPE_CREDENTIALS,
            data=structs_pb2.AuthDataClientUnion(credentials=structs_pb2.Credentials(
                type=enums_pb2.CREDENTIALS_TYPE_MFILES, username=username, password=password)))
        # [END login-password]
        return self._log_in(authentication)

    def log_in_with_token(self, token: str, plugin: str, scope: str,
                          plugin_vault: str = "") -> structs_pb2.SessionData:
        """Logs in with a token from the vault's identity provider (SSO).

        The token goes to the authentication plugin as the single named value ``Token``.
        ``plugin`` and ``scope`` are the plugin's name and its system configuration value
        ``Scope`` from GetAuthenticationConfiguration; ``plugin_vault`` is its value
        ``VaultGUID``, when present.
        """
        # [START login-token]
        plugin_data = structs_pb2.AuthDataClientUnion.PluginData(
            plugin_name=plugin,
            configuration_scope=scope,
            vault_guid=plugin_vault,
            authentication_attempt_identifier=structs_pb2.AuthenticationAttemptIdentifier(
                data=bytes(32)),
            data_format=enums_pb2.PLUGIN_AUTH_DATA_FORMAT_UNENCRYPTED,
            auth_data=[structs_pb2.NamedValue(key="Token", value=structs_pb2.TypedValue(
                type=enums_pb2.DATATYPE_TEXT, data=structs_pb2.TypedValueUnion(text=token)))])
        authentication = structs_pb2.AuthDataClient(
            type=enums_pb2.AUTH_DATA_TYPE_PLUGIN,
            data=structs_pb2.AuthDataClientUnion(plugin=plugin_data))
        # [END login-token]
        return self._log_in(authentication)

    def _log_in(self, authentication: structs_pb2.AuthDataClient) -> structs_pb2.SessionData:
        # [START login]
        request = IRPCLogin_pb2.LogInRequest(
            environment_data=structs_pb2.EnvironmentData(
                client_name=self.client_name, host_platform="Python"),
            client_data=self._client_data(),
            login_data=structs_pb2.LoginData(
                server_hostname=self.config.host, vault_guid=self.config.vault),
            authentication_data=authentication)
        response = self.login.LogIn(request)
        if not response.session_data.session_id:
            raise RuntimeError("LogIn did not complete: the authentication needs more steps")
        # From now on, the interceptor sends this in "mfiles-session-id-bin" on every call.
        self.session_id = response.session_data.session_id
        # [END login]
        self.session = response.session_data
        return self.session

    def switch_to_vault(self, vault_guid: str) -> structs_pb2.SessionData:
        """From a server-level session, opens a session to the vault without authenticating again,
        then ends the server-level session."""
        # [START login-with-session]
        response = self.login.LogInWithSession(IRPCLogin_pb2.LogInWithSessionRequest(
            client_data=self._client_data(),
            login_data=structs_pb2.LoginData(server_hostname=self.config.host,
                                             vault_guid=vault_guid)))
        self.log_out()  # The server-level session; the new vault session stays valid.
        self.config.vault = vault_guid
        self.session_id = response.session_data.session_id
        self.session = response.session_data
        # [END login-with-session]
        return self.session

    def _client_data(self) -> structs_pb2.ClientData:
        return structs_pb2.ClientData(
            type=enums_pb2.CLIENT_TYPE_SERVER_API, language="en",
            capabilities=structs_pb2.ClientCapabilities(long_authentication_attempt_identifier=True))

    def log_out(self) -> None:
        if self.session_id:
            try:
                self.login.LogOut(IRPCLogin_pb2.LogOutRequest())
            finally:
                self.session_id = b""
                self.session = None

    def close(self) -> None:
        """Logs out (if logged in) and closes the channel."""
        try:
            self.log_out()
        except grpc.RpcError:
            pass  # The session may already have expired; the server cleans it up.
        finally:
            self._channel.close()

    def __enter__(self) -> "MFilesClient":
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    def _record_metadata(self, headers, trailers) -> None:
        """Keeps the cookies and the metadata structure version that a call returned."""
        for key, value in [*(headers or ()), *(trailers or ())]:
            if key == SET_COOKIE_KEY:
                # M-Files Cloud pins a client to one server instance with the cookie mfilesmsm.
                # Keep its name and value and send it back on every later call.
                name, _, value = value.split(";", 1)[0].strip().partition("=")
                if name:
                    self.cookies[name] = value
            elif key == METADATA_VERSION_KEY and len(value) == 4:
                # The vault's metadata structure version: a little-endian 32-bit integer.
                self.metadata_structure_version = struct.unpack("<i", value)[0]


def decode_version(version: int) -> str:
    """Decodes a version from QueryVersion: four 16-bit parts, e.g. "26.11.16500.0"."""
    return ".".join(str((version >> shift) & 0xFFFF) for shift in (48, 32, 16, 0))


def run(main) -> None:
    """Runs a sample's main function; prints a failed call's M-Files error and exits with 1."""
    try:
        main()
    except MFilesError as error:
        print(error, file=sys.stderr)
        sys.exit(1)