# SPDX-FileCopyrightText: Christian Amsüss and the aiocoap contributors
#
# SPDX-License-Identifier: MIT

"""a plain CoAP proxy that can work both as forward and as reverse proxy"""

import sys
import argparse

import aiocoap
from aiocoap.proxy.server import (
    ForwardProxyWithPooledObservations,
    ProxyWithPooledObservations,
    NameBasedVirtualHost,
    SubdomainVirtualHost,
    SubresourceVirtualHost,
    UnconditionalRedirector,
)
from aiocoap.util.cli import AsyncCLIDaemon
from aiocoap.cli.common import add_server_arguments, server_context_from_arguments


def build_parser():
    p = argparse.ArgumentParser(description=__doc__)

    mode = p.add_argument_group(
        "mode", "Required argument for setting the operation mode"
    )
    mode.add_argument("--forward", help="Run as forward proxy", action="store_true")
    mode.add_argument("--reverse", help="Run as reverse proxy", action="store_true")

    details = p.add_argument_group(
        "details", "Options that govern how requests go in and out"
    )
    add_server_arguments(details)
    details.add_argument(
        "--register",
        help="Register with a Resource directory",
        metavar="RD-URI",
        nargs="?",
        default=False,
    )
    details.add_argument(
        "--register-as",
        help="Endpoint name (with possibly a domain after a dot) to register as",
        metavar="EP[.D]",
        default=None,
    )
    details.add_argument(
        "--register-proxy",
        help="Ask the RD to serve as a reverse proxy. Note that this is only practical for --unconditional or --pathbased reverse proxies.",
        action="store_true",
    )

    r = p.add_argument_group(
        "Rules",
        description="Sequence of forwarding rules "
        "that, if matched by a request, specify a forwarding destination. Destinations can be prefixed to change their behavior: With an '@' sign, they are treated as forward proxies. With a '!' sign, the destination is set as Uri-Host.",
    )

    class TypedAppend(argparse.Action):
        def __call__(self, parser, namespace, values, option_string=None):
            if getattr(namespace, self.dest) is None:
                setattr(namespace, self.dest, [])
            getattr(namespace, self.dest).append((option_string, values))

    r.add_argument(
        "--namebased",
        help="If Uri-Host matches NAME, route to DEST",
        metavar="NAME:DEST",
        action=TypedAppend,
        dest="r",
    )
    r.add_argument(
        "--subdomainbased",
        help="If Uri-Host is anything.NAME, route to DEST",
        metavar="NAME:DEST",
        action=TypedAppend,
        dest="r",
    )
    r.add_argument(
        "--pathbased",
        help="If a requested path starts with PATH, split that part off and route to DEST",
        metavar="PATH:DEST",
        action=TypedAppend,
        dest="r",
    )
    r.add_argument(
        "--unconditional",
        help="Route all requests not previously matched to DEST",
        metavar="DEST",
        action=TypedAppend,
        dest="r",
    )

    return p


def destsplit(dest):
    use_as_proxy = False
    rewrite_uri_host = False
    if dest.startswith("!"):
        dest = dest[1:]
        rewrite_uri_host = True
    if dest.startswith("@"):
        dest = dest[1:]
        use_as_proxy = True
    return dest, rewrite_uri_host, use_as_proxy


class Main(AsyncCLIDaemon):
    async def start(self, args=None):
        parser = build_parser()
        options = parser.parse_args(args if args is not None else sys.argv[1:])
        self.options = options

        if not options.forward and not options.reverse:
            raise parser.error("At least one of --forward and --reverse must be given.")

        self.outgoing_context = await aiocoap.Context.create_client_context()
        if options.forward:
            proxy = ForwardProxyWithPooledObservations(self.outgoing_context)
        else:
            proxy = ProxyWithPooledObservations(self.outgoing_context)
        for kind, data in options.r or ():
            if kind in ("--namebased", "--subdomainbased"):
                try:
                    name, dest = data.split(":", 1)
                except Exception:
                    raise parser.error("%s needs NAME:DEST as arguments" % kind)
                dest, rewrite_uri_host, use_as_proxy = destsplit(dest)
                if rewrite_uri_host and kind == "--subdomainbased":
                    parser.error(
                        "The flag '!' makes no sense for subdomain based redirection as the subdomain data would be lost"
                    )
                r = (
                    NameBasedVirtualHost
                    if kind == "--namebased"
                    else SubdomainVirtualHost
                )(name, dest, rewrite_uri_host, use_as_proxy)
            elif kind == "--pathbased":
                try:
                    path, dest = data.split(":", 1)
                except Exception:
                    raise parser.error("--pathbased needs PATH:DEST as arguments")
                r = SubresourceVirtualHost(path.split("/"), dest)
            elif kind == "--unconditional":
                dest, rewrite_uri_host, use_as_proxy = destsplit(data)
                if rewrite_uri_host:
                    parser.error(
                        "The flag '!' makes no sense for unconditional redirection as the host name data would be lost"
                    )
                r = UnconditionalRedirector(dest, use_as_proxy)
            else:
                raise AssertionError("Unknown redirectory kind")
            proxy.add_redirector(r)

        self.proxy_context = await server_context_from_arguments(proxy, options)

        if options.register is not False:
            from aiocoap.resourcedirectory.client.register import Registerer

            params = {}
            if options.register_as:
                ep, _, d = options.register_as.partition(".")
                params["ep"] = ep
                if d:
                    params["d"] = d
            if options.register_proxy:
                # FIXME: Check this in discovery
                params["proxy"] = "on"
            # FIXME: Construct this from settings (path-based), and forward results
            proxy.get_resources_as_linkheader = lambda: ""
            self.registerer = Registerer(
                self.proxy_context,
                rd=options.register,
                lt=60,
                registration_parameters=params,
            )

    async def shutdown(self):
        if self.options.register is not False:
            await self.registerer.shutdown()
        await self.outgoing_context.shutdown()
        await self.proxy_context.shutdown()


sync_main = Main.sync_main

if __name__ == "__main__":
    # if you want to run this using `python3 -m`, see http://bugs.python.org/issue22480
    sync_main()