# PubSub::Subscribe

// PubSub::Subscribe — unary Pull of a Google Cloud Pub/Sub subscription.
//
// Volang speaks unary HTTP and has no bidirectional stream, so this uses the
// REST Pull API rather than StreamingPull. One request asks for one message.
// returnImmediately keeps an empty subscription from holding the only HTTP
// slot: the server answers at once, and the block waits poll_interval.
// A message is acknowledged before it is written to the outputs. The next
// pull is then queued immediately, so a backlog does not wait out poll_interval.
// An empty response leaves the scheduler to ask again after that interval.
//
// Phases: 0 idle, 1 token, 2 pull, 3 acknowledge.

fn build_jwt() {
    sa_key = config::get("service_account_key")
    private_key = json::get(sa_key, "private_key")
    client_email = json::get(sa_key, "client_email")
    key_id = json::get(sa_key, "private_key_id")

    now = time::now()
    exp = now + 3600

    header = str::fmt("""{"alg":"RS256","typ":"JWT","kid":"{}"}""", key_id)
    claims = str::fmt("""{"iss":"{}","sub":"{}","scope":"https://www.googleapis.com/auth/pubsub","aud":"https://oauth2.googleapis.com/token","iat":{},"exp":{}}""", client_email, client_email, now, exp)

    header_b64 = base64::url_encode(header)
    claims_b64 = base64::url_encode(claims)
    signing_input = str::fmt("{}.{}", header_b64, claims_b64)
    signature = crypto::rs256_sign(signing_input, private_key)

    return str::fmt("{}.{}", signing_input, signature)
}

fn configured() {
    if (config::get("service_account_key") == "") {
        return false
    }
    if (config::get("project_id") == "") {
        return false
    }
    return config::get("subscription") != ""
}

// json::get aborts when the path is missing, so probe the raw text first.
fn has_field(body, key) {
    return str::count(body, str::fmt("""{}":""", key)) > 0
}

fn subscription_url(method) {
    return str::fmt("https://pubsub.googleapis.com/v1/projects/{}/subscriptions/{}:{}", config::get("project_id"), config::get("subscription"), method)
}

fn fail(status) {
    state::set("phase", 0)
    state::set("pending_ack_id", "")
    state::set("pending_message", "")
    output::set("status", status)
    output::set("success", false)
    if (status == 401) {
        state::set("token_expiry", 0)
    }
}

// A request left in flight (a dropped response) must not wedge the block.
fn recover_if_stuck() {
    if (state::get("phase") == 0) {
        return false
    }
    if (time::uptime() - state::get("phase_started") <= 10000) {
        return true
    }
    fail(0)
    return false
}

fn request_token() {
    jwt = build_jwt()

    client = http::client()
    http::set_method(client, "POST")
    http::set_header(client, "Content-Type", "application/x-www-form-urlencoded")
    http::set_body(client, str::fmt("grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer&assertion={}", jwt))
    http::call(client, "https://oauth2.googleapis.com/token")

    state::set("phase", 1)
    state::set("phase_started", time::uptime())
}

fn do_pull() {
    client = http::client()
    http::set_method(client, "POST")
    http::set_header(client, "Content-Type", "application/json")
    http::set_header(client, "Authorization", str::fmt("Bearer {}", state::get("token")))
    http::set_body(client, """{"maxMessages":1,"returnImmediately":true}""")
    http::call(client, subscription_url("pull"))

    state::set("phase", 2)
    state::set("phase_started", time::uptime())
}

fn do_ack() {
    client = http::client()
    http::set_method(client, "POST")
    http::set_header(client, "Content-Type", "application/json")
    http::set_header(client, "Authorization", str::fmt("Bearer {}", state::get("token")))
    http::set_body(client, str::fmt("""{"ackIds":["{}"]}""", state::get("pending_ack_id")))
    http::call(client, subscription_url("acknowledge"))

    state::set("phase", 3)
    state::set("phase_started", time::uptime())
}

fn start() {
    if (!configured()) {
        return
    }
    if (recover_if_stuck()) {
        return
    }
    if (state::get("token_expiry") > time::uptime()) {
        do_pull()
    } else {
        request_token()
    }
}

// The scheduler's next slot was already chosen when this pull started, and it
// will not run twice in the same second. A callback is what drains the next
// message without waiting out poll_interval.
fn delivered() {
    state::set("phase", 0)
    state::set("pending_ack_id", "")
    state::set("pending_message", "")
    callback::clear()
    callback::set(0, "onPull", 0)
}

extern fn onPull(value) {
    start()
}

extern fn http::on_response(status, body, headers) {
    phase = state::get("phase")

    if (phase == 1) {
        if (status == 200) {
            state::set("token", json::get(body, "access_token"))
            state::set("token_expiry", time::uptime() + 3500000)
            do_pull()
        } else {
            fail(status)
        }
        return
    }

    if (phase == 2) {
        if (status != 200) {
            fail(status)
            return
        }
        if (!has_field(body, "receivedMessages")) {
            state::set("phase", 0)
            output::set("status", status)
            output::set("success", true)
            return
        }
        messages = json::get(body, "receivedMessages")
        if (array::len(messages) == 0) {
            state::set("phase", 0)
            output::set("status", status)
            output::set("success", true)
            return
        }

        item = array::get_at(messages, 0)
        if (!has_field(item, "ackId")) {
            state::set("phase", 0)
            output::set("status", status)
            output::set("success", false)
            return
        }

        state::set("pending_ack_id", json::get(item, "ackId"))
        if (has_field(item, "data")) {
            encoded = json::get(item, "message.data")
            if (encoded == "") {
                state::set("pending_message", "")
            } else {
                state::set("pending_message", base64::decode(encoded))
            }
        } else {
            state::set("pending_message", "")
        }
        // Acknowledge before the outputs change, so a failed ack leaves the
        // message unreleased and it is delivered again.
        do_ack()
        return
    }

    if (phase == 3) {
        if (status == 200) {
            output::set("status", status)
            output::set("success", true)
            output::set("message", state::get("pending_message"))
            output::set("seq", output::get("seq") + 1)
            delivered()
        } else {
            fail(status)
        }
        return
    }
}

extern fn std::next_activation_at() {
    if (state::get("last_activation_ts") == 0) {
        return time::now()
    }
    interval = math::int(math::round(config::get("poll_interval")))
    if (interval < 1) {
        interval = 1
    }
    return time::now() + interval
}

channel = input::channel()

if (channel == "config_change") {
    state::set("token_expiry", 0)
}

if (channel == "" or channel == "config_change") {
    start()
}