import { AthenaClient, StartQueryExecutionCommand, GetQueryExecutionCommand, GetQueryResultsCommand } from "@aws-sdk/client-athena";
const athenaClient = new AthenaClient({ region: process.env.AWS_REGION, profile: 'default' });

const executionTimeout = 1000;

const productQuery = `SELECT DISTINCT
	line_item_resource_id AS resourceId,
	line_item_product_code AS productCode,
	line_item_usage_account_id AS accountId
FROM ${process.env.ATHENA_CU_TABLE};`;

export const executeQueryAsync = async (sqlQuery) => {
    const startQueryCommand = new StartQueryExecutionCommand({
        QueryString: sqlQuery,
        QueryExecutionContext: { Database: process.env.ATHENA_CU_DATABASE },
        ResultConfiguration: { OutputLocation: process.env.ATHENA_OUTPUT_S3_BUCKET },
    });

    const startQueryResponse = await athenaClient.send(startQueryCommand);
    const queryExecutionId = startQueryResponse.QueryExecutionId;

    let queryExecutionStatus;
    do {
        const getQueryExecutionCommand = new GetQueryExecutionCommand({
            QueryExecutionId: queryExecutionId,
        });
        const queryExecutionResponse = await athenaClient.send(getQueryExecutionCommand);
        queryExecutionStatus = queryExecutionResponse.QueryExecution.Status.State;

        if (queryExecutionStatus === "FAILED") {
            console.error(`Query Failed: ${JSON.stringify(queryExecutionResponse)}`);
            break;
        }
        await new Promise((resolve) => setTimeout(resolve, executionTimeout));
    } while (queryExecutionStatus !== "SUCCEEDED");

    return queryExecutionId;
}

export const retrieveResultsAsync = async (queryExecutionId) => {
    const getQueryResultsCommand = new GetQueryResultsCommand({
        QueryExecutionId: queryExecutionId,
    });

    const result = await athenaClient.send(getQueryResultsCommand);
    if (!result || !result.ResultSet || !result.ResultSet.Rows) {
        return {
            statusCode: 404,
            body: JSON.stringify([]),
        };
    }
    let rows = result.ResultSet.Rows;
    if (0 == rows.length || 1 == rows.length) {
        return {
            statusCode: 400,
            body: JSON.stringify([]),
        };
    }
    const columnNames = rows[0].Data.map(item => item.VarCharValue);
    const items = [];
    rows.slice(1).forEach(cells => {
        const item = {};
        for (let i = 0; i < cells.Data.length; i++) {
            item[columnNames[i]] = cells.Data[i].VarCharValue;
        }
        items.push(item);
    });
    return items;
};