# **************************************************************************
#                                  _   _ ____  _
#  Project                     ___| | | |  _ \| |
#                             / __| | | | |_) | |
#                            | (__| |_| |  _ <| |___
#                             \___|\___/|_| \_\_____|
#
# Copyright (C) Daniel Stenberg, <daniel@haxx.se>, et al.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at https://curl.se/docs/copyright.html.
#
# You may opt to use, copy, modify, merge, publish, distribute and/or sell
# copies of the Software, and permit persons to whom the Software is
# furnished to do so, under the terms of the COPYING file.
#
# This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY
# KIND, either express or implied.
#
# SPDX-License-Identifier: curl
#
# **************************************************************************

"""
HTTP/2 server push
"""

import sys
import ctypes as ct
from dataclasses import dataclass
from pathlib import Path

import libcurl as lcurl
from curl_utils import *  # noqa
from debug import debug_function

here = Path(__file__).resolve().parent

OUT_DIR  = here/"output"
OUT_FILE = OUT_DIR/"dl"


@dataclass
class debug_config:
    trace_ascii: bool = False
    outstream: object = None


@lcurl.write_callback
def write_function(buffer, size, nitems, userp):
    config = lcurl.from_oid(userp)
    buffer_size = nitems * size
    if buffer_size == 0: return 0
    bwritten = bytes(buffer[:buffer_size])
    nwritten = config.outstream.write(bwritten)
    return nwritten


count = 0

@lcurl.push_callback
def server_push_callback(parent, easy, num_headers, headers, userp):
    # called when there is an incoming push

    transfersp = ct.cast(userp, ct.POINTER(ct.c_int))

    global count

    file_path = OUT_DIR/("push%u" % count)
    count += 1

    # here's a new stream, save it in a new file for each new push
    try:
        #FILE *out;
        out = file_path.open("wb")
    except:
        # if we cannot save it, deny it
        print("Failed to create output file for push", file=sys.stderr)
        return lcurl.CURL_PUSH_DENY

    # write to this file
    #!!!lcurl.easy_setopt(easy, lcurl.CURLOPT_WRITEDATA, out)

    print("**** push callback approves stream %d, got %d headers!" %
          (count, num_headers), file=sys.stderr)

    for i in range(num_headers):
        headp = lcurl.pushheader_bynum(headers, i)
        print("**** header %d: %s" % (i, headp.encode("utf-8")),
              file=sys.stderr)

    headp = lcurl.pushheader_byname(headers, b":path")
    if headp:
        print("**** The PATH is %s" % headp.encode("utf-8"),  # skip :path + colon
              file=sys.stderr)

    transfersp.contents += 1  # one more

    return lcurl.CURL_PUSH_OK


def setup(curl: ct.POINTER(lcurl.CURL), config: debug_config, url: str) -> int:

    global OUT_FILE

    try:
        config.outstream = OUT_FILE.open("wb")
    except OSError as exc:
        config.outstream = None
        print("error: could not open file %s for writing: %s" %
              (OUT_FILE, exc.strerror), file=sys.stderr)
        return 1  # failed

    # set the URL
    lcurl.easy_setopt(curl, lcurl.CURLOPT_URL, url.encode("utf-8"))
    # send all data to this function
    lcurl.easy_setopt(curl, lcurl.CURLOPT_WRITEFUNCTION, write_function)
    # write to this file
    lcurl.easy_setopt(curl, lcurl.CURLOPT_WRITEDATA, id(config))
    # please be verbose
    lcurl.easy_setopt(curl, lcurl.CURLOPT_VERBOSE, 1)
    lcurl.easy_setopt(curl, lcurl.CURLOPT_DEBUGFUNCTION, debug_function)
    lcurl.easy_setopt(curl, lcurl.CURLOPT_DEBUGDATA, id(config))
    # HTTP/2 please
    lcurl.easy_setopt(curl, lcurl.CURLOPT_HTTP_VERSION,
                            lcurl.CURL_HTTP_VERSION_2_0)
    # we use a self-signed test server, skip verification during debugging
    lcurl.easy_setopt(curl, lcurl.CURLOPT_SSL_VERIFYPEER, 0)
    lcurl.easy_setopt(curl, lcurl.CURLOPT_SSL_VERIFYHOST, 0)
    if lcurl.CURLPIPE_MULTIPLEX > 0:
        # wait for pipe connection to confirm
        lcurl.easy_setopt(curl, lcurl.CURLOPT_PIPEWAIT, 1)

    return 0  # all is good


#
# Download a file over HTTP/2, take care of server push.
#

def main(argv=sys.argv[1:]):

    url: str = argv[0] if len(argv) >= 1 else "https://localhost:8443/index.html"

    config = debug_config(True)  # enable ascii tracing

    # init a multi stack
    mcurl: ct.POINTER(lcurl.CURLM) = lcurl.multi_init()
    curl:  ct.POINTER(lcurl.CURL)  = lcurl.easy_init()

    with curl_guard(False, None, mcurl) as guard:
        if not curl: return 1

        # set options
        res = setup(curl, config, url)
        if res:
            print("failed", file=sys.stderr)
            return res

        # add the easy transfer
        lcurl.multi_add_handle(mcurl, curl)

        transfers = ct.c_int(1)  # we start with one
        lcurl.multi_setopt(mcurl, lcurl.CURLMOPT_PIPELINING,
                                  lcurl.CURLPIPE_MULTIPLEX)
        lcurl.multi_setopt(mcurl, lcurl.CURLMOPT_PUSHFUNCTION,
                                  server_push_callback)
        lcurl.multi_setopt(mcurl, lcurl.CURLMOPT_PUSHDATA,
                                  ct.byref(transfers))

        while transfers.value:  # as long as we have transfers going

            still_running = ct.c_int(1)  # keep number of running handles
            mcode: int = lcurl.multi_perform(mcurl, ct.byref(still_running))
            # wait for activity, timeout or "nothing"
            if still_running.value: mcode = lcurl.multi_poll(mcurl, None, 0, 1000, None)
            if mcode:
                break

            # A little caution when doing server push is that libcurl itself
            # has created and added one or more easy handles but we need to
            # clean them up when we are done.
            while True:
                msgs_left = ct.c_int(0)
                msgp: ct.POINTER(lcurl.CURLMsg) = lcurl.multi_info_read(mcurl,
                                                                        ct.byref(msgs_left))
                if not msgp: break
                msg = msgp.contents

                if msg.msg == lcurl.CURLMSG_DONE:
                    transfers.value -= 1
                    lcurl.multi_remove_handle(mcurl, msg.easy_handle)
                    lcurl.easy_cleanup(msg.easy_handle)

    return 0


sys.exit(main())
