#include "include.h"

#include <access/xact.h>
#include <miscadmin.h>
#include <pgstat.h>
#include <replication/slot.h>
#include <storage/proc.h>
#include <tcop/tcopprot.h>
#include <unistd.h>
#include <utils/builtins.h>
#include <utils/lsyscache.h>
#include <utils/memutils.h>
#include <utils/ps_status.h>
#include <utils/snapmgr.h>
#include <utils/timeout.h>

#if PG_VERSION_NUM >= 110000
#include <jit/jit.h>
#endif

static Task task = {0};

Task *get_task(void) {
    return &task;
}

static char *SPI_getvalue_my(TupleTableSlot *slot, TupleDesc tupdesc, int fnumber) {
    bool isnull;
    bool typisvarlena;
    Datum attr = slot_getattr(slot, fnumber, &isnull);
    Oid foutoid;
    if (isnull) return NULL;
    getTypeOutputInfo(TupleDescAttr(tupdesc, fnumber - 1)->atttypid, &foutoid, &typisvarlena);
    return OidOutputFunctionCall(foutoid, attr);
}

static void headers(TupleDesc tupdesc) {
    if (task.output.len) appendStringInfoString(&task.output, "\n");
    for (int col = 1; col <= tupdesc->natts; col++) {
        char *fname = SPI_fname(tupdesc, col);
        if (col > 1) appendStringInfoChar(&task.output, task.delimiter);
        appendBinaryStringInfoEscapeQuote(&task.output, fname, strlen(fname), false, task.escape, task.quote);
        pfree(fname);
    }
}

static
#if PG_VERSION_NUM >= 90600
bool
#else
void
#endif
receiveSlot(TupleTableSlot *slot, DestReceiver *self) {
    TupleDesc tupdesc = slot->tts_tupleDescriptor;
    if (!task.shared)
        return
#if PG_VERSION_NUM >= 90600
    true
#endif
    ;
    if (!task.output.data) initStringInfoMy(&task.output);
    if (task.header && !task.row && tupdesc->natts > 1) headers(tupdesc);
    if (task.output.len) appendStringInfoString(&task.output, "\n");
    for (int col = 1; col <= tupdesc->natts; col++) {
        char *value = SPI_getvalue_my(slot, tupdesc, col);
        if (col > 1) appendStringInfoChar(&task.output, task.delimiter);
        if (!value) appendStringInfoString(&task.output, task.null); else {
            appendBinaryStringInfoEscapeQuote(&task.output, value, strlen(value), !init_oid_is_string(SPI_gettypeid(tupdesc, col)) && task.string, task.escape, task.quote);
            pfree(value);
        }
    }
    task.row++;
#if PG_VERSION_NUM >= 90600
    return true;
#endif
}

static void rStartup(DestReceiver *self, int operation, TupleDesc tupdesc) {
    if (!task.shared) return;
    switch (operation) {
        case CMD_UNKNOWN: elog(DEBUG1, "id = %li, operation = CMD_UNKNOWN", task.shared->id); break;
        case CMD_SELECT: elog(DEBUG1, "id = %li, operation = CMD_SELECT", task.shared->id); break;
        case CMD_UPDATE: elog(DEBUG1, "id = %li, operation = CMD_UPDATE", task.shared->id); break;
        case CMD_INSERT: elog(DEBUG1, "id = %li, operation = CMD_INSERT", task.shared->id); break;
        case CMD_DELETE: elog(DEBUG1, "id = %li, operation = CMD_DELETE", task.shared->id); break;
        case CMD_UTILITY: elog(DEBUG1, "id = %li, operation = CMD_UTILITY", task.shared->id); break;
        case CMD_NOTHING: elog(DEBUG1, "id = %li, operation = CMD_NOTHING", task.shared->id); break;
        default: elog(DEBUG1, "id = %li, operation = %i", task.shared->id, operation); break;
    }
    task.row = 0;
    task.skip = operation == CMD_SELECT;
}

static void rShutdown(DestReceiver *self) {
    if (task.shared) elog(DEBUG1, "id = %li", task.shared->id);
}

static void rDestroy(DestReceiver *self) {
    if (task.shared) elog(DEBUG1, "id = %li", task.shared->id);
}

static
#if PG_VERSION_NUM >= 120000
const
#endif
DestReceiver myDestReceiver = {
    .receiveSlot = receiveSlot,
    .rStartup = rStartup,
    .rShutdown = rShutdown,
    .rDestroy = rDestroy,
    .mydest = DestDebug,
};

DestReceiver *CreateDestReceiverMy(CommandDest dest) {
    if (task.shared) elog(DEBUG1, "id = %li", task.shared->id);
#if PG_VERSION_NUM >= 120000
    return unconstify(DestReceiver *, &myDestReceiver);
#else
    return &myDestReceiver;
#endif
}

static void ReadyForQueryMy(CommandDest dest) {
    if (task.shared) elog(DEBUG1, "id = %li", task.shared->id);
}

void NullCommandMy(CommandDest dest) {
    if (task.shared) elog(DEBUG1, "id = %li", task.shared->id);
}

#if PG_VERSION_NUM >= 130000
void BeginCommandMy(CommandTag commandTag, CommandDest dest) {
    if (task.shared) elog(DEBUG1, "id = %li, commandTag = %s", task.shared->id, GetCommandTagName(commandTag));
}

void EndCommandMy(const QueryCompletion *qc, CommandDest dest, bool force_undecorated_output) {
    char completionTag[COMPLETION_TAG_BUFSIZE];
    CommandTag tag = qc->commandTag;
    const char *tagname = GetCommandTagName(tag);
    if (!task.shared) return;
    if (command_tag_display_rowcount(tag) && !force_undecorated_output) snprintf(completionTag, COMPLETION_TAG_BUFSIZE, tag == CMDTAG_INSERT ? "%s 0 %lu" : "%s %lu", tagname, qc->nprocessed);
    else snprintf(completionTag, COMPLETION_TAG_BUFSIZE, "%s", tagname);
    elog(DEBUG1, "id = %li, completionTag = %s", task.shared->id, completionTag);
    if (task.skip) task.skip = 0; else {
        if (!task.output.data) initStringInfoMy(&task.output);
        if (task.output.len) appendStringInfoString(&task.output, "\n");
        appendStringInfoString(&task.output, completionTag);
    }
}
#else
void BeginCommandMy(const char *commandTag, CommandDest dest) {
    if (task.shared) elog(DEBUG1, "id = %li, commandTag = %s", task.shared->id, commandTag);
}

void EndCommandMy(const char *commandTag, CommandDest dest) {
    if (!task.shared) return;
    elog(DEBUG1, "id = %li, commandTag = %s", task.shared->id, commandTag);
    if (task.skip) task.skip = 0; else {
        if (!task.output.data) initStringInfoMy(&task.output);
        if (task.output.len) appendStringInfoString(&task.output, "\n");
        appendStringInfoString(&task.output, commandTag);
    }
}
#endif

static void dest_execute_spi(const char *src) {
    bool count = false;
    bool insert = false;
    char completionTag[COMPLETION_TAG_BUFSIZE];
    int rc = SPI_execute(src, false, 0);
    const char *tagname = SPI_result_code_string(rc) + (rc >= 0 ? sizeof("SPI_OK_") - 1 : sizeof("SPI_ERROR_") - 1);
    switch (rc) {
        case SPI_ERROR_ARGUMENT: ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), errmsg("invalid arguments"))); break;
        case SPI_ERROR_COPY: ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), errmsg("COPY is not supported"))); break;
        case SPI_ERROR_OPUNKNOWN: ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), errmsg("unrecognized command type"))); break;
        case SPI_ERROR_TRANSACTION: ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), errmsg("transaction control statement is not supported"))); break;
        case SPI_OK_DELETE: count = true; break;
        case SPI_OK_DELETE_RETURNING: count = true; break;
        case SPI_OK_INSERT: count = true; insert = true; break;
        case SPI_OK_INSERT_RETURNING: count = true; insert = true; break;
        case SPI_OK_SELECT: count = true; task.skip = 1; break;
        case SPI_OK_UPDATE: count = true; break;
        case SPI_OK_UPDATE_RETURNING: count = true; break;
    }
    elog(DEBUG1, "id = %li, commandTag = %s", task.shared->id, tagname);
    if (SPI_tuptable) for (uint64 row = 0; row < SPI_processed; row++) {
        task.skip = 1;
        if (!task.output.data) initStringInfoMy(&task.output);
        if (task.header && !row && SPI_tuptable->tupdesc->natts > 1) headers(SPI_tuptable->tupdesc);
        if (task.output.len) appendStringInfoString(&task.output, "\n");
        for (int col = 1; col <= SPI_tuptable->tupdesc->natts; col++) {
            char *value = SPI_getvalue(SPI_tuptable->vals[row], SPI_tuptable->tupdesc, col);
            if (col > 1) appendStringInfoChar(&task.output, task.delimiter);
            if (!value) appendStringInfoString(&task.output, task.null); else {
                appendBinaryStringInfoEscapeQuote(&task.output, value, strlen(value), !init_oid_is_string(SPI_gettypeid(SPI_tuptable->tupdesc, col)) && task.string, task.escape, task.quote);
                pfree(value);
            }
        }
    }
    if (count) snprintf(completionTag, COMPLETION_TAG_BUFSIZE, insert ? "%s 0 %lu" : "%s %lu", tagname, (unsigned long)SPI_processed);
    else snprintf(completionTag, COMPLETION_TAG_BUFSIZE, "%s", tagname);
    elog(DEBUG1, "id = %li, completionTag = %s", task.shared->id, completionTag);
    if (task.skip) task.skip = 0; else {
        if (!task.output.data) initStringInfoMy(&task.output);
        if (task.output.len) appendStringInfoString(&task.output, "\n");
        appendStringInfoString(&task.output, completionTag);
    }
}

static void dest_execute(void) {
    if (!task.shared->spi) {
        ListCell *cell;
        MemoryContext oldMemoryContext = MemoryContextSwitchTo(MessageContext);
        MemoryContextResetAndDeleteChildren(MessageContext);
        InvalidateCatalogSnapshotConditionally();
        foreach(cell, pg_parse_query(task.input)) {
#if PG_VERSION_NUM >= 100000
            Node *node = ((RawStmt *)lfirst(cell))->stmt;
#else
            Node *node = (Node *)lfirst(cell);
#endif
            if (IsA(node, CopyStmt) && !((CopyStmt *)node)->filename) ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), errmsg("COPY %s is not supported", ((CopyStmt *)node)->is_from ? "FROM STDIN" : "TO STDOUT")));
        }
        MemoryContextSwitchTo(oldMemoryContext);
        whereToSendOutput = DestDebug;
        ReadyForQueryMy(whereToSendOutput);
        SetCurrentStatementStartTimestamp();
        exec_simple_query_my(task.input);
        if (IsTransactionState()) {
            exec_simple_query_my(SQL(END));
            if (IsTransactionState()) ereport(ERROR, (errcode(ERRCODE_ACTIVE_SQL_TRANSACTION), errmsg("still active sql transaction")));
        }
    } else {
#if PG_VERSION_NUM >= 100000
        // RawStmt.stmt_location lets us slice task.input into the text of each individual
        // statement and run them through SPI one by one, the same way exec_simple_query and
        // the remote libpq path already report a result per statement instead of just the last.
        List *parsetree_list = pg_parse_query(task.input);
        if (list_length(parsetree_list) <= 1) dest_execute_spi(task.input); else {
            ListCell *cell;
            int prev_start = -1;
            foreach(cell, parsetree_list) {
                int start = lfirst_node(RawStmt, cell)->stmt_location;
                if (prev_start >= 0) {
                    char *stmt = pnstrdup(task.input + prev_start, start - prev_start);
                    dest_execute_spi(stmt);
                    pfree(stmt);
                }
                prev_start = start;
            }
            dest_execute_spi(task.input + prev_start);
        }
#else
        dest_execute_spi(task.input);
#endif
    }
}

static void dest_catch(void) {
    if (!task.shared->spi) {
        HOLD_INTERRUPTS();
        disable_all_timeouts(false);
        QueryCancelPending = false;
    }
    EmitErrorReport();
    if (!task.shared->spi) {
        debug_query_string = NULL;
        AbortOutOfAnyTransaction();
#if PG_VERSION_NUM >= 110000
        PortalErrorCleanup();
#endif
        if (MyReplicationSlot) ReplicationSlotRelease();
#if PG_VERSION_NUM >= 170000
        ReplicationSlotCleanup(false);
#elif PG_VERSION_NUM >= 100000
        ReplicationSlotCleanup();
#endif
#if PG_VERSION_NUM >= 110000
        jit_reset_after_error();
#endif
        MemoryContextSwitchTo(TopMemoryContext);
    }
    FlushErrorState();
    if (!task.shared->spi) {
        xact_started_my(false);
        RESUME_INTERRUPTS();
    }
}

static void dest_role(const char *role, bool reset) {
    if (task.shared->spi) SetConfigOption("role", reset ? "none" : role, PGC_USERSET, PGC_S_SESSION); else {
        Shared *shared = task.shared;
        StringInfoData src;
        task.shared = NULL; // disable dest receiver and command tags while switching role
        initStringInfoMy(&src);
        if (reset) appendStringInfoString(&src, SQL(RESET ROLE;)); else {
            const char *quote = quote_identifier(role);
            appendStringInfo(&src, SQL(SET ROLE %s;), quote);
            if (quote != role) pfree((void *)quote);
        }
        PG_TRY();
            exec_simple_query_my(src.data);
        PG_CATCH();
            task.shared = shared; // restore before any error handling dereferences it
            PG_RE_THROW();
        PG_END_TRY();
        pfree(src.data);
        task.shared = shared;
    }
}

static void dest_discard(void) {
    Shared *shared = task.shared;
    static const char *src = SQL(SET SESSION AUTHORIZATION DEFAULT; RESET ALL; DEALLOCATE ALL; CLOSE ALL; UNLISTEN *; DISCARD PLANS; DISCARD TEMP; DISCARD SEQUENCES;);
    StringInfoData oid;
    task.shared = NULL; // disable dest receiver and command tags during cleanup
    PG_TRY();
        if (shared->spi) {
            SPI_connect_my(src);
            SPI_execute_with_args_my(src, 0, NULL, NULL, NULL, SPI_OK_UTILITY);
            SPI_finish_my();
        } else exec_simple_query_my(src);
    PG_CATCH();
        task.shared = shared; // restore before any error handling dereferences it
        PG_RE_THROW();
    PG_END_TRY();
    task.shared = shared;
    SetConfigOption("search_path", "", PGC_USERSET, PGC_S_SESSION);
    SetConfigOption("pg_task.schema", task.shared->schema, PGC_USERSET, PGC_S_SESSION);
    SetConfigOption("pg_task.table", task.shared->table, PGC_USERSET, PGC_S_SESSION);
    initStringInfoMy(&oid);
    appendStringInfo(&oid, "%i", task.shared->oid);
    SetConfigOption("pg_task.oid", oid.data, PGC_USERSET, PGC_S_SESSION);
    pfree(oid.data);
}

bool dest_timeout(void) {
    bool exit;
    int StatementTimeoutMy = StatementTimeout;
    if (task_work(&task)) return true;
    elog(DEBUG1, "id = %li, timeout = %i, input = %s, count = %i", task.shared->id, task.timeout, task.input, task.count);
    set_ps_display_my("timeout");
    StatementTimeout = task.timeout;
    if (task.shared->spi) {
        SPI_connect_my(task.input);
        BeginInternalSubTransaction(NULL);
    }
    PG_TRY();
        SetConfigOption("search_path", task_search_path(), PGC_USERSET, PGC_S_SESSION);
        dest_role(task.user, false);
        dest_execute();
        dest_role(NULL, true);
        SetConfigOption("search_path", "", PGC_USERSET, PGC_S_SESSION);
        if (task.shared->spi) ReleaseCurrentSubTransaction();
    PG_CATCH();
        SetConfigOption("search_path", "", PGC_USERSET, PGC_S_SESSION);
        task_error(&task);
        dest_catch();
        if (task.shared->spi) {
            RollbackAndReleaseCurrentSubTransaction();
#if PG_VERSION_NUM < 100000
            SPI_restore_connection();
#endif
        }
        dest_role(NULL, true);
    PG_END_TRY();
    if (task.shared->spi) SPI_finish_my();
    StatementTimeout = StatementTimeoutMy;
    pgstat_report_stat(false);
    pgstat_report_activity(STATE_IDLE, NULL);
    set_ps_display_my("idle");
    exit = task_done(&task, true);
    if (!exit && !task.save) dest_discard();
    return exit;
}