/*
 * Copyright (c) 2021 Fastly, Kazuho Oku
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */
#include <inttypes.h>
#include "h2o.h"

#define SEND_WAIT 500        /* how frequent new loglines should be pushed out; in milliseconds */
#define BUFFER_LIMIT 8388608 /* do not buffer more than 8MB */

struct h2o_self_trace_generator {
    h2o_generator_t super;
    h2o_req_t *req;
    h2o_buffer_t *buf;
    h2o_doublebuffer_t inflight;
    h2o_timer_t send_timer;
    unsigned should_send_buffered : 1;
};

static void on_generator_dispose(void *_self)
{
    struct h2o_self_trace_generator *self = _self;

    quicly_tracer_t *tracer = self->req->conn->callbacks->get_tracer(self->req->conn);
    tracer->cb = NULL;
    tracer->ctx = NULL;

    h2o_buffer_dispose(&self->buf);
    h2o_doublebuffer_dispose(&self->inflight);
    h2o_timer_unlink(&self->send_timer);
}

static void do_send(struct h2o_self_trace_generator *self)
{
    h2o_iovec_t vec = h2o_doublebuffer_prepare(&self->inflight, &self->buf, SIZE_MAX);
    h2o_send(self->req, &vec, 1, H2O_SEND_STATE_IN_PROGRESS);
}

static void on_send_timeout(h2o_timer_t *_timer)
{
    struct h2o_self_trace_generator *self = H2O_STRUCT_FROM_MEMBER(struct h2o_self_trace_generator, send_timer, _timer);

    assert(!self->inflight.inflight);
    assert(self->should_send_buffered);

    self->should_send_buffered = 0;
    do_send(self);
}

static void adjust_send_timer(struct h2o_self_trace_generator *self)
{
    if (!self->inflight.inflight && self->should_send_buffered && !h2o_timer_is_linked(&self->send_timer))
        h2o_timer_link(self->req->conn->ctx->loop, SEND_WAIT, &self->send_timer);
}

static void log_trace(void *_self, const char *fmt, ...)
{
    struct h2o_self_trace_generator *self = _self;

    /* append provided input to the buffer */
    if (self->buf->size < BUFFER_LIMIT) {
        va_list args;
        va_start(args, fmt);

        h2o_iovec_t buf = h2o_buffer_reserve(&self->buf, 1024);
        int len = vsnprintf(buf.base, buf.len, fmt, args);
        if (len >= buf.len) {
            buf = h2o_buffer_reserve(&self->buf, len + 1);
            len = vsnprintf(buf.base, buf.len, fmt, args);
            assert(len < buf.len);
        }
        self->buf->size += len;

        va_end(args);
    }

    /* Log is sent only when there's another request inflight. Otherwise, it is buffered until another request becomes inflight. */
    if (!self->should_send_buffered && self->req->conn->callbacks->num_reqs_inflight(self->req->conn) > 1)
        self->should_send_buffered = 1;
    adjust_send_timer(self);
}

static void do_proceed(h2o_generator_t *_self, h2o_req_t *_req)
{
    struct h2o_self_trace_generator *self = (void *)_self;

    assert(self->inflight.inflight);
    h2o_doublebuffer_consume(&self->inflight);

    adjust_send_timer(self);
}

static int on_req(h2o_handler_t *handler, h2o_req_t *req)
{
    if (req->conn->callbacks->get_tracer == NULL) {
        h2o_send_error_403(req, "Forbidden", "not available", 0);
        return 0;
    }

    quicly_tracer_t *tracer = req->conn->callbacks->get_tracer(req->conn);
    if (tracer->cb != NULL) {
        h2o_send_error_403(req, "Forbidden", "conn-state handler is already attached", 0);
        return 0;
    }

    /* instantiate the generator */
    struct h2o_self_trace_generator *self = h2o_mem_alloc_shared(&req->pool, sizeof(*self), on_generator_dispose);
    self->super.proceed = do_proceed;
    self->super.stop = NULL;
    self->req = req;
    h2o_buffer_init(&self->buf, &h2o_socket_buffer_prototype);
    h2o_doublebuffer_init(&self->inflight, &h2o_socket_buffer_prototype);
    self->send_timer = (h2o_timer_t){.cb = on_send_timeout};
    self->should_send_buffered = 0;

    /* register */
    tracer->cb = log_trace;
    tracer->ctx = self;

    /* build response headers */
    req->res.status = 200;
    h2o_add_header(&req->pool, &req->res.headers, H2O_TOKEN_CONTENT_TYPE, NULL, H2O_STRLIT("text/plain; charset=utf-8"));
    h2o_add_header(&req->pool, &req->res.headers, H2O_TOKEN_CACHE_CONTROL, NULL, H2O_STRLIT("no-cache, no-store"));
    h2o_buffer_append(&self->buf, "\n", 1); /* add some data for simplicity */

    h2o_start_response(self->req, &self->super);
    do_send(self);

    return 0;
}

void h2o_self_trace_register(h2o_pathconf_t *conf)
{
    h2o_handler_t *self = h2o_create_handler(conf, sizeof(*self));
    self->on_req = on_req;
}
