"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.captureAWSClient = void 0;
const service_error_classification_1 = require("@smithy/service-error-classification");
const aws_1 = __importDefault(require("../segments/attributes/aws"));
const querystring_1 = require("querystring");
const subsegment_1 = __importDefault(require("../segments/attributes/subsegment"));
const contextUtils = require('../context_utils');
const logger = require('../logger');
const { safeParseInt } = require('../utils');
const utils_1 = require("../utils");
const XRAY_PLUGIN_NAME = 'XRaySDKInstrumentation';
const buildAttributesFromMetadata = async (service, operation, region, commandInput, res, error) => {
    var _a, _b, _c;
    const { extendedRequestId, requestId, httpStatusCode: statusCode, attempts } = ((_a = res === null || res === void 0 ? void 0 : res.output) === null || _a === void 0 ? void 0 : _a.$metadata) || (error === null || error === void 0 ? void 0 : error.$metadata);
    const aws = new aws_1.default({
        extendedRequestId,
        requestId,
        retryCount: attempts,
        data: res === null || res === void 0 ? void 0 : res.output,
        request: {
            operation,
            params: commandInput,
            httpRequest: {
                region,
                statusCode,
            },
        },
    }, service);
    const http = {};
    if (statusCode) {
        http.response = {};
        http.response.status = statusCode;
    }
    if (((_b = res === null || res === void 0 ? void 0 : res.response) === null || _b === void 0 ? void 0 : _b.headers) && ((_c = res === null || res === void 0 ? void 0 : res.response) === null || _c === void 0 ? void 0 : _c.headers['content-length']) !== undefined) {
        if (!http.response) {
            http.response = {};
        }
        http.response.content_length = safeParseInt(res.response.headers['content-length']);
    }
    return [aws, http];
};
function addFlags(http, subsegment, err) {
    var _a, _b, _c;
    if (err && (0, service_error_classification_1.isThrottlingError)(err)) {
        subsegment.addThrottleFlag();
    }
    else if (safeParseInt((_a = http.response) === null || _a === void 0 ? void 0 : _a.status) === 429 || safeParseInt((_b = err === null || err === void 0 ? void 0 : err.$metadata) === null || _b === void 0 ? void 0 : _b.httpStatusCode) === 429) {
        subsegment.addThrottleFlag();
    }
    const cause = (0, utils_1.getCauseTypeFromHttpStatus)(safeParseInt((_c = http.response) === null || _c === void 0 ? void 0 : _c.status));
    if (cause === 'fault') {
        subsegment.addFaultFlag();
    }
    else if (cause === 'error') {
        subsegment.addErrorFlag();
    }
}
const getXRayMiddleware = (config, manualSegment) => (next, context) => async (args) => {
    var _a;
    const segment = contextUtils.isAutomaticMode() ? contextUtils.resolveSegment() : manualSegment;
    const { clientName, commandName } = context;
    const commandInput = (_a = args === null || args === void 0 ? void 0 : args.input) !== null && _a !== void 0 ? _a : {};
    const commandOperation = commandName.slice(0, -7); // Strip trailing "Command" string
    const operation = commandOperation.charAt(0).toLowerCase() + commandOperation.slice(1);
    const service = clientName.slice(0, -6); // Strip trailing "Client" string
    if (!segment) {
        const output = service + '.' + operation;
        if (!contextUtils.isAutomaticMode()) {
            logger.getLogger().info('Call ' + output + ' requires a segment object' +
                ' passed to captureAWSv3Client for tracing in manual mode. Ignoring.');
        }
        else {
            logger.getLogger().info('Call ' + output +
                ' is missing the sub/segment context for automatic mode. Ignoring.');
        }
        return next(args);
    }
    let subsegment;
    if (segment.notTraced) {
        subsegment = segment.addNewSubsegmentWithoutSampling(service);
    }
    else {
        subsegment = segment.addNewSubsegment(service);
    }
    subsegment.addAttribute('namespace', 'aws');
    const parent = (segment instanceof subsegment_1.default ? segment.segment : segment);
    const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData;
    let traceHeader = (0, querystring_1.stringify)({
        Root: parent.trace_id,
        Parent: subsegment.id,
        Sampled: subsegment.notTraced ? '0' : '1',
    }, ';');
    if (data != null) {
        for (const [key, value] of Object.entries(data)) {
            traceHeader += ';' + key + '=' + value;
        }
    }
    if (!segment.noOp) {
        args.request.headers['X-Amzn-Trace-Id'] = traceHeader;
    }
    let res;
    try {
        res = await next(args);
        if (!res) {
            throw new Error('Failed to get response from instrumented AWS Client.');
        }
        const [aws, http] = await buildAttributesFromMetadata(service, operation, await config.region(), commandInput, res, null);
        subsegment.addAttribute('aws', aws);
        subsegment.addAttribute('http', http);
        addFlags(http, subsegment);
        subsegment.close();
        return res;
    }
    catch (err) {
        if (err.$metadata) {
            const [aws, http] = await buildAttributesFromMetadata(service, operation, await config.region(), commandInput, null, err);
            subsegment.addAttribute('aws', aws);
            subsegment.addAttribute('http', http);
            addFlags(http, subsegment, err);
        }
        const errObj = { message: err.message, name: err.name, stack: err.stack || new Error().stack };
        subsegment.close(errObj, true);
        throw err;
    }
};
const xRayMiddlewareOptions = {
    name: XRAY_PLUGIN_NAME,
    step: 'build',
};
const getXRayPlugin = (config, manualSegment) => ({
    applyToStack: (stack) => {
        stack.add(getXRayMiddleware(config, manualSegment), xRayMiddlewareOptions);
    },
});
/**
 * Instruments AWS SDK V3 clients with X-Ray via middleware.
 *
 * @param client - AWS SDK V3 client to instrument
 * @param manualSegment - Parent segment or subsegment that is passed in for manual mode users
 * @returns - the client with the X-Ray instrumentation middleware added to its middleware stack
 */
function captureAWSClient(client, manualSegment) {
    // Remove existing middleware to ensure operation is idempotent
    client.middlewareStack.remove(XRAY_PLUGIN_NAME);
    client.middlewareStack.use(getXRayPlugin(client.config, manualSegment));
    return client;
}
exports.captureAWSClient = captureAWSClient;
