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

"""aiocoap-client is a simple command-line tool for interacting with CoAP servers"""

import sys
import asyncio
import argparse
import logging
import subprocess
from pathlib import Path

import shlex

# even though not used directly, this has side effects on the input() function
# used in interactive mode
try:
    import readline  # noqa: F401
except ImportError:
    pass  # that's normal on some platforms, and ok since it's just a usability enhancement

import aiocoap
import aiocoap.defaults
import aiocoap.meta
import aiocoap.proxy.client
from aiocoap.util import contenttype
from aiocoap.util.cli import ActionNoYes
from aiocoap.numbers import ContentFormat


def build_parser():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument(
        "--non",
        help="Send request as non-confirmable (NON) message",
        action="store_true",
    )
    p.add_argument(
        "-m",
        "--method",
        help="Name or number of request method to use (default: %(default)s)",
        default="GET",
    )
    p.add_argument(
        "--observe", help="Register an observation on the resource", action="store_true"
    )
    p.add_argument(
        "--observe-exec",
        help="Run the specified program whenever the observed resource changes, feeding the response data to its stdin",
        metavar="CMD",
    )
    p.add_argument(
        "--accept",
        help="Content format to request",
        metavar="MIME",
    )
    p.add_argument(
        "--proxy", help="Relay the CoAP request to a proxy for execution", metavar="URI"
    )
    p.add_argument(
        "--payload",
        help="Send X as request payload (eg. with a PUT). If X starts with an '@', its remainder is treated as a file name and read from; '@-' reads from the console. Non-file data may be recoded, see --content-format.",
        metavar="X",
    )
    p.add_argument(
        "--payload-initial-szx",
        help="Size exponent to limit the initial block's size (0 ≙ 16 Byte, 6 ≙ 1024 Byte)",
        metavar="SZX",
        type=int,
    )
    p.add_argument(
        "--content-format",
        help="Content format of the --payload data. If a known format is given and --payload has a non-file argument, the payload is converted from CBOR Diagnostic Notation.",
        metavar="MIME",
    )
    p.add_argument(
        "--no-set-hostname",
        help="Suppress transmission of Uri-Host even if the host name is not an IP literal",
        dest="set_hostname",
        action="store_false",
        default=True,
    )
    p.add_argument(
        "-v",
        "--verbose",
        help="Increase the debug output",
        action="count",
    )
    p.add_argument(
        "-q",
        "--quiet",
        help="Decrease the debug output",
        action="count",
    )
    # careful: picked before parsing
    p.add_argument(
        "--interactive",
        help="Enter interactive mode",
        action="store_true",
    )
    p.add_argument(
        "--credentials",
        help="Load credentials to use from a given file",
        type=Path,
    )
    p.add_argument(
        "--version", action="version", version="%(prog)s " + aiocoap.meta.version
    )

    p.add_argument(
        "--color",
        help="Color output (default on TTYs if all required modules are installed)",
        default=None,
        action=ActionNoYes,
    )
    p.add_argument(
        "--pretty-print",
        help="Pretty-print known content formats (default on TTYs if all required modules are installed)",
        default=None,
        action=ActionNoYes,
    )
    p.add_argument(
        "url",
        help="CoAP address to fetch",
    )

    return p


def configure_logging(verbosity):
    logging.basicConfig()

    if verbosity <= -2:
        logging.getLogger("coap").setLevel(logging.CRITICAL + 1)
    elif verbosity == -1:
        logging.getLogger("coap").setLevel(logging.ERROR)
    elif verbosity == 0:
        logging.getLogger("coap").setLevel(logging.WARNING)
    elif verbosity == 1:
        logging.getLogger("coap").setLevel(logging.INFO)
    elif verbosity >= 2:
        logging.getLogger("coap").setLevel(logging.DEBUG)


def colored(text, options, tokenlambda):
    """Apply pygments based coloring if options.color is set. Tokelambda is a
    callback to which pygments.token is passed and which returns a token type;
    this makes it easy to not need to conditionally react to pygments' possible
    absence in all color locations."""
    if not options.color:
        return str(text)

    from pygments.formatters import TerminalFormatter
    from pygments import token, format

    return format(
        [(tokenlambda(token), str(text))],
        TerminalFormatter(),
    )


def incoming_observation(options, response):
    if options.observe_exec:
        p = subprocess.Popen(options.observe_exec, shell=True, stdin=subprocess.PIPE)
        # FIXME this blocks
        p.communicate(response.payload)
    else:
        sys.stdout.write(colored("---", options, lambda token: token.Comment.Preproc))
        if response.code.is_successful():
            present(response, options, file=sys.stderr)
        else:
            sys.stdout.flush()
            print(
                colored(
                    response.code, options, lambda token: token.Token.Generic.Error
                ),
                file=sys.stderr,
            )
            if response.payload:
                present(response, options, file=sys.stderr)


def apply_credentials(context, credentials, errfn):
    if credentials.suffix == ".json":
        import json

        context.client_credentials.load_from_dict(json.load(credentials.open("rb")))
    elif credentials.suffix == ".diag":
        try:
            import cbor_diag
            import cbor2
        except ImportError:
            raise errfn(
                "Loading credentials in CBOR diagnostic format requires cbor2 and cbor_diag package"
            )
        context.client_credentials.load_from_dict(
            cbor2.loads(cbor_diag.diag2cbor(credentials.open().read()))
        )
    else:
        raise errfn(
            "Unknown suffix: %s (expected: .json or .diag)" % (credentials.suffix)
        )


def present(message, options, file=sys.stdout):
    """Write a message payload to the output, pretty printing and/or coloring
    it as configured in the options."""
    if not options.quiet and (message.opt.location_path or message.opt.location_query):
        # FIXME: Percent encoding is completely missing; this would be done
        # most easily with a CRI library
        location_ref = "/" + "/".join(message.opt.location_path)
        if message.opt.location_query:
            location_ref += "?" + "&".join(message.opt.location_query)
        print(
            colored(
                f"Location options indicate new resource: {location_ref}",
                options,
                lambda token: token.Token.Generic.Inserted,
            ),
            file=sys.stderr,
        )

    if not message.payload:
        return

    payload = None

    cf = message.opt.content_format or message.request.opt.content_format
    if cf is not None and cf.is_known():
        mime = cf.media_type
    else:
        mime = "application/octet-stream"
    if options.pretty_print:
        from aiocoap.util.prettyprint import pretty_print

        prettyprinted = pretty_print(message)
        if prettyprinted is not None:
            (infos, mime, payload) = prettyprinted
            if not options.quiet:
                for i in infos:
                    print(
                        colored("# " + i, options, lambda token: token.Comment),
                        file=sys.stderr,
                    )

    color = options.color
    if color:
        from aiocoap.util.prettyprint import lexer_for_mime
        import pygments

        try:
            lexer = lexer_for_mime(mime)
        except pygments.util.ClassNotFound:
            color = False

    if color and payload is None:
        # Coloring requires a unicode-string style payload, either from the
        # mime type or from the pretty printer.
        try:
            payload = message.payload.decode("utf8")
        except UnicodeDecodeError:
            color = False

    if color:
        from pygments.formatters import TerminalFormatter
        from pygments import highlight

        highlit = highlight(
            payload,
            lexer,
            TerminalFormatter(),
        )
        # The TerminalFormatter already adds an end-of-line character, not
        # trying to add one for any missing trailing newlines.
        print(highlit, file=file, end="")
        file.flush()
    else:
        if payload is None:
            file.buffer.write(message.payload)
            if file.isatty() and message.payload[-1:] != b"\n":
                file.write("\n")
        else:
            file.write(payload)
            if file.isatty() and payload[-1] != "\n":
                file.write("\n")


async def single_request(args, context):
    parser = build_parser()
    options = parser.parse_args(args)

    pretty_print_modules = aiocoap.defaults.prettyprint_missing_modules()
    if pretty_print_modules and (options.color is True or options.pretty_print is True):
        parser.error(
            "Color and pretty printing require the following"
            " additional module(s) to be installed: %s"
            % ", ".join(pretty_print_modules)
        )
    if options.color is None:
        options.color = sys.stdout.isatty() and not pretty_print_modules
    if options.pretty_print is None:
        options.pretty_print = sys.stdout.isatty() and not pretty_print_modules

    configure_logging((options.verbose or 0) - (options.quiet or 0))

    try:
        try:
            code = getattr(
                aiocoap.numbers.codes.Code,
                options.method.upper().replace("IPATCH", "iPATCH"),
            )
        except AttributeError:
            try:
                code = aiocoap.numbers.codes.Code(int(options.method))
            except ValueError:
                raise parser.error("Unknown method")

        if options.credentials is not None:
            apply_credentials(context, options.credentials, parser.error)

        request = aiocoap.Message(
            code=code, mtype=aiocoap.NON if options.non else aiocoap.CON
        )
        request.set_request_uri(options.url, set_uri_host=options.set_hostname)

        if options.accept:
            try:
                request.opt.accept = ContentFormat(int(options.accept))
            except ValueError:
                try:
                    request.opt.accept = ContentFormat.by_media_type(options.accept)
                except KeyError:
                    raise parser.error("Unknown accept type")

        if options.observe:
            request.opt.observe = 0
            observation_is_over = asyncio.get_event_loop().create_future()

        if options.content_format:
            try:
                request.opt.content_format = ContentFormat(int(options.content_format))
            except ValueError:
                try:
                    request.opt.content_format = ContentFormat.by_media_type(
                        options.content_format
                    )
                except KeyError:
                    raise parser.error("Unknown content format")

        if options.payload:
            if options.payload.startswith("@"):
                filename = options.payload[1:]
                if filename == "-":
                    f = sys.stdin.buffer
                else:
                    f = open(filename, "rb")
                try:
                    request.payload = f.read()
                except OSError as e:
                    raise parser.error("File could not be opened: %s" % e)
            else:
                request_classification = contenttype.categorize(
                    request.opt.content_format.media_type
                    if request.opt.content_format is not None
                    and request.opt.content_format.is_known()
                    else ""
                )
                if request_classification in ("cbor", "cbor-seq"):
                    try:
                        import cbor_diag
                    except ImportError as e:
                        raise parser.error(f"CBOR recoding not available ({e})")

                    try:
                        encoded = cbor_diag.diag2cbor(options.payload)
                    except ValueError as e:
                        raise parser.error(
                            f"Parsing CBOR diagnostic notation failed. Make sure quotation marks are escaped from the shell. Error: {e}"
                        )

                    if request_classification == "cbor-seq":
                        try:
                            import cbor2
                        except ImportError as e:
                            raise parser.error(
                                f"CBOR sequence recoding not available ({e})"
                            )
                        decoded = cbor2.loads(encoded)
                        if not isinstance(decoded, list):
                            raise parser.error(
                                "CBOR sequence recoding requires an array as the top-level element."
                            )
                        request.payload = b"".join(cbor2.dumps(d) for d in decoded)
                    else:
                        request.payload = encoded
                else:
                    request.payload = options.payload.encode("utf8")

        if options.payload_initial_szx is not None:
            request.remote.maximum_block_size_exp = options.payload_initial_szx

        if options.proxy is None:
            interface = context
        else:
            interface = aiocoap.proxy.client.ProxyForwarder(options.proxy, context)

        requested_uri = request.get_request_uri()

        requester = interface.request(request)

        if options.observe:
            requester.observation.register_errback(observation_is_over.set_result)
            requester.observation.register_callback(
                lambda data, options=options: incoming_observation(options, data)
            )

        try:
            response_data = await requester.response
        finally:
            if not requester.response.done():
                requester.response.cancel()
            if options.observe and not requester.observation.cancelled:
                requester.observation.cancel()

        response_uri = response_data.get_request_uri()
        if requested_uri != response_uri:
            print(
                colored(
                    f"Response arrived from different address; base URI is {response_uri}",
                    options,
                    lambda token: token.Generic.Inserted,
                ),
                file=sys.stderr,
            )
        if response_data.code.is_successful():
            present(response_data, options)
        else:
            print(
                colored(response_data.code, options, lambda token: token.Generic.Error),
                file=sys.stderr,
            )
            present(response_data, options, file=sys.stderr)
            sys.exit(1)

        if options.observe:
            exit_reason = await observation_is_over
            print("Observation is over: %r" % (exit_reason,), file=sys.stderr)
    except aiocoap.error.HelpfulError as e:
        print(str(e), file=sys.stderr)
        extra_help = e.extra_help(
            hints=dict(
                original_uri=options.url,
                request=request,
            )
        )
        if extra_help:
            print("Debugging hint:", extra_help, file=sys.stderr)
        sys.exit(1)
    # Fallback while not all backends raise NetworkErrors
    except OSError as e:
        text = str(e)
        if not text:
            text = repr(e)
        if not text:
            # eg ConnectionResetError flying out of a misconfigured SSL server
            text = type(e)
        print(
            "Warning: OS errors should not be raised this way any more.",
            file=sys.stderr,
        )
        # not telling what to do precisely: the form already tells users to
        # include `aiocoap.cli.defaults` output, which is exactly what we
        # need.
        print(
            f"Even if the cause of the error itself is clear, please file an issue at {aiocoap.meta.bugreport_uri}.",
            file=sys.stderr,
        )
        print("Error:", text, file=sys.stderr)
        sys.exit(1)


async def single_request_with_context(args):
    """Wrapper around single_request until sync_main gets made fully async, and
    async context managers are used to manage contexts."""
    context = await aiocoap.Context.create_client_context()
    try:
        await single_request(args, context)
    finally:
        await context.shutdown()


interactive_expecting_keyboard_interrupt = None


async def interactive():
    global interactive_expecting_keyboard_interrupt
    interactive_expecting_keyboard_interrupt = asyncio.get_event_loop().create_future()

    context = await aiocoap.Context.create_client_context()

    while True:
        try:
            # when http://bugs.python.org/issue22412 is resolved, use that instead
            line = await asyncio.get_event_loop().run_in_executor(
                None, lambda: input("aiocoap> ")
            )
        except EOFError:
            line = "exit"
        line = shlex.split(line)
        if not line:
            continue
        if line in (["help"], ["?"]):
            line = ["--help"]
        if line in (["quit"], ["q"], ["exit"]):
            break

        current_task = asyncio.create_task(
            single_request(line, context=context),
            name="Interactive prompt command %r" % line,
        )
        interactive_expecting_keyboard_interrupt = (
            asyncio.get_event_loop().create_future()
        )

        done, pending = await asyncio.wait(
            [current_task, interactive_expecting_keyboard_interrupt],
            return_when=asyncio.FIRST_COMPLETED,
        )

        if current_task not in done:
            current_task.cancel()
        else:
            try:
                await current_task
            except SystemExit as e:
                if e.code != 0:
                    print("Exit code: %d" % e.code, file=sys.stderr)
                continue
            except Exception as e:
                print("Unhandled exception raised: %s" % (e,))

    await context.shutdown()


def sync_main(args=None):
    # interactive mode is a little messy, that's why this is not using aiocoap.util.cli yet
    if args is None:
        args = sys.argv[1:]

    if "--interactive" not in args:
        try:
            asyncio.run(single_request_with_context(args))
        except KeyboardInterrupt:
            sys.exit(3)
    else:
        if len(args) != 1:
            print(
                "No other arguments must be specified when entering interactive mode",
                file=sys.stderr,
            )
            sys.exit(1)

        loop = asyncio.get_event_loop()
        task = loop.create_task(
            interactive(),
            name="Interactive prompt",
        )

        while not task.done():
            try:
                loop.run_until_complete(task)
            except KeyboardInterrupt:
                if not interactive_expecting_keyboard_interrupt.done():
                    interactive_expecting_keyboard_interrupt.set_result(None)
            except SystemExit:
                continue  # asyncio/tasks.py(242) raises those after setting them as results, but we particularly want them back in the loop


if __name__ == "__main__":
    sync_main()
