import { encodableSchema } from "@automate.ax/codec"
import { Client, type QueryResult } from "pg"
import { parse } from "pgsql-parser"
import * as z from "zod"
import { defineAction, terminalActionError } from "../../automation/actions"
import { defineAccount } from "../../automation/integrations"

const POSTGRES_SECRET_SCHEMA = z.object({
  connectionString: z.string().min(1),
})
const POSTGRES_PARAMETER_SCHEMA = z.union([
  z.json(),
  z.bigint(),
  z.date(),
  z.instanceof(Uint8Array),
])
const POSTGRES_STATEMENT_SCHEMA = z.object({
  /** Values bound to `$1`, `$2`, and subsequent placeholders. */
  parameters: POSTGRES_PARAMETER_SCHEMA.array().optional(),

  /** One SQL statement to execute. */
  sql: z.string().trim().min(1),
})
const POSTGRES_FIELD_SCHEMA = z.object({
  /** One-based source column number, or zero for a computed field. */
  columnId: z.number().int(),

  /** PostgreSQL data type OID. */
  dataTypeId: z.number().int(),

  /** Type modifier supplied by PostgreSQL. */
  dataTypeModifier: z.number().int(),

  /** Storage size of the PostgreSQL data type. */
  dataTypeSize: z.number().int(),

  /** Wire format used for the field. */
  format: z.string(),

  /** Column name returned by PostgreSQL. */
  name: z.string(),

  /** Source table OID, or zero for a computed field. */
  tableId: z.number().int(),
})
const POSTGRES_ROW_SCHEMA = z.record(z.string(), encodableSchema)
const POSTGRES_RESULT_SCHEMA = z.object({
  /** PostgreSQL command tag such as SELECT, INSERT, or UPDATE. */
  command: z.string(),

  /** Metadata for each returned field. */
  fields: POSTGRES_FIELD_SCHEMA.array(),

  /** Affected or returned row count when PostgreSQL supplies one. */
  rowCount: z.number().int().nonnegative().nullable(),

  /** Result rows keyed by returned column name. */
  rows: POSTGRES_ROW_SCHEMA.array(),
})

/** Declares a named PostgreSQL account binding for an automation. */
export const postgresAccount = defineAccount("postgres")

/**
 * Runs one parameterized statement in a provider-enforced read-only
 * transaction.
 *
 * A fresh connection is used for each action invocation. Results retain
 * provider command and field metadata alongside the returned rows.
 */
export const runPostgresQuery = defineAction("Run PostgreSQL query")
  .describe(
    "Runs one parameterized read-only statement and returns rows plus PostgreSQL result metadata.",
  )
  .account("postgres")
  .input(POSTGRES_STATEMENT_SCHEMA)
  .output(POSTGRES_RESULT_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(async ({ account, input }) => {
    await validatePostgresStatement(input.sql)

    return withPostgresClient(account.secret, async (client) => {
      await client.query("BEGIN READ ONLY")
      try {
        // Validate the public output before committing so schema failures roll back.
        const normalized = normalizePostgresResult(
          await client.query({
            text: input.sql,
            values: input.parameters,
          }),
        )
        await client.query("COMMIT")
        return normalized
      } catch (error) {
        await rollbackPostgresTransaction(client, error)
        throw error
      }
    })
  })

/**
 * Executes one parameterized PostgreSQL statement in its own transaction.
 *
 * Use `runPostgresQuery` for provider-enforced read-only work and this action
 * for statements that can change data or database objects.
 */
export const executePostgresStatement = defineAction(
  "Execute PostgreSQL statement",
)
  .describe(
    "Executes one parameterized read/write statement and returns rows plus PostgreSQL result metadata.",
  )
  .account("postgres")
  .input(POSTGRES_STATEMENT_SCHEMA)
  .output(POSTGRES_RESULT_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(async ({ account, input }) => {
    await validatePostgresStatement(input.sql)

    return withPostgresClient(account.secret, async (client) =>
      normalizePostgresResult(
        await client.query({ text: input.sql, values: input.parameters }),
      ),
    )
  })

/**
 * Executes ordered parameterized statements in one PostgreSQL transaction.
 *
 * The transaction rolls back after the first statement failure. Transaction
 * control statements are rejected because this action owns the boundary.
 */
export const runPostgresTransaction = defineAction("Run PostgreSQL transaction")
  .describe(
    "Executes ordered parameterized statements atomically and returns one result for each statement.",
  )
  .account("postgres")
  .input(
    z.object({
      /** Ordered statements to execute atomically. */
      statements: POSTGRES_STATEMENT_SCHEMA.array().min(1).max(100),
    }),
  )
  .output(z.object({ results: POSTGRES_RESULT_SCHEMA.array() }))
  .retry({ replaySafety: "unsafe" })
  .handler(async ({ account, input }) => {
    await Promise.all(
      input.statements.map(({ sql }) => validatePostgresStatement(sql)),
    )

    return withPostgresClient(account.secret, async (client) => {
      await client.query("BEGIN")
      try {
        const results = []
        for (const statement of input.statements) {
          results.push(
            normalizePostgresResult(
              await client.query({
                text: statement.sql,
                values: statement.parameters,
              }),
            ),
          )
        }
        await client.query("COMMIT")
        return { results }
      } catch (error) {
        await rollbackPostgresTransaction(client, error)
        throw error
      }
    })
  })

/** PostgreSQL callables. */
export const postgres = {
  execute: executePostgresStatement,
  query: runPostgresQuery,
  transaction: runPostgresTransaction,
}

/**
 * Creates a PostgreSQL client from a resolved integration secret.
 *
 * Custom actions can use this helper when they need COPY, cursors, prepared
 * session reuse, advisory locks, or another provider feature beyond the
 * packaged actions.
 *
 * @param secret - Validated PostgreSQL integration secret.
 */
export function getPostgresClient(secret: Record<string, unknown>) {
  return new Client(POSTGRES_SECRET_SCHEMA.parse(secret))
}

/**
 * Verifies that packaged actions can safely own the statement boundary.
 *
 * @param sql - SQL source to validate.
 */
async function validatePostgresStatement(sql: string) {
  let statements
  try {
    statements = (await parse(sql)).stmts ?? []
  } catch (error) {
    throw terminalActionError(
      error instanceof Error
        ? error
        : new Error("PostgreSQL could not parse the SQL statement.", {
            cause: error,
          }),
    )
  }
  if (statements.length !== 1) {
    throw terminalActionError(
      new Error("PostgreSQL actions accept exactly one SQL statement."),
    )
  }

  const statement = statements[0]
  if (statement?.stmt && "TransactionStmt" in statement.stmt) {
    throw terminalActionError(
      new Error(
        "PostgreSQL actions own their transaction boundaries and do not accept transaction-control statements.",
      ),
    )
  }
}

/**
 * Converts node-postgres metadata to the public camel-case result contract.
 *
 * @param result - Raw node-postgres result.
 */
function normalizePostgresResult(result: QueryResult) {
  return POSTGRES_RESULT_SCHEMA.parse({
    command: result.command,
    fields: result.fields.map((field) => ({
      columnId: field.columnID,
      dataTypeId: field.dataTypeID,
      dataTypeModifier: field.dataTypeModifier,
      dataTypeSize: field.dataTypeSize,
      format: field.format,
      name: field.name,
      tableId: field.tableID,
    })),
    rowCount: result.rowCount,
    rows: result.rows,
  })
}

/**
 * Rolls back a failed transaction while retaining both provider failures.
 *
 * @param client - Connected PostgreSQL client.
 * @param cause - Original statement failure.
 */
async function rollbackPostgresTransaction(client: Client, cause: unknown) {
  let rollbackFailure: unknown
  try {
    await client.query("ROLLBACK")
    return
  } catch (error) {
    rollbackFailure = error
  }

  const rollbackRecord =
    typeof rollbackFailure === "object" && rollbackFailure !== null
      ? rollbackFailure
      : undefined
  throw new Error(
    "PostgreSQL statement failed and its transaction could not be rolled back.",
    {
      cause: Object.assign(
        new Error(
          rollbackFailure instanceof Error
            ? rollbackFailure.message
            : String(rollbackFailure),
          { cause },
        ),
        rollbackRecord && "code" in rollbackRecord
          ? { code: rollbackRecord.code }
          : {},
      ),
    },
  )
}

/**
 * Runs work on a fresh PostgreSQL connection and always closes it.
 *
 * @param secret - Validated PostgreSQL integration secret.
 * @param execute - Work to run on the connected client.
 */
async function withPostgresClient<T>(
  secret: Record<string, unknown>,
  execute: (client: Client) => Promise<T>,
) {
  const client = getPostgresClient(secret)
  try {
    await client.connect()
    return await execute(client)
  } finally {
    await client.end()
  }
}