import * as z from "zod"
import { defineAction } from "../../../automation/actions"
import {
  buildA1Range,
  getGoogleSheetsApi,
  normalizeRowValues,
  normalizeSpreadsheetId,
  readSourceRows,
  resolveRowSource,
  toGoogleCellData,
} from "../lib/google-sheets"
import { GOOGLE_SHEETS_WRITE_REQUIREMENT } from "../lib/scopes"
import {
  ROW_SOURCE_SCHEMA,
  ROW_VALUES_INPUT_SCHEMA,
  SPREADSHEET_REFERENCE_SCHEMA,
  VALUE_INPUT_MODE_SCHEMA,
} from "../lib/schemas"

/** Updates rows matching key columns and appends records with new keys. */
export const upsertGoogleSheetsRows = defineAction("Upsert Google Sheets rows")
  .describe("Updates unique matching rows and appends records not yet present.")
  .account("google", GOOGLE_SHEETS_WRITE_REQUIREMENT)
  .input(
    z.object({
      /** Exact headers whose combined values uniquely identify a row. */
      keyColumns: z.string().min(1).array().min(1),

      /** Whether strings are parsed like UI input or stored literally. */
      inputMode: VALUE_INPUT_MODE_SCHEMA.optional(),

      /** One record or a non-empty list of records. */
      rows: ROW_VALUES_INPUT_SCHEMA,

      /** Header-backed sheet or native table. */
      source: ROW_SOURCE_SCHEMA,

      /** Spreadsheet ID or standard Google Sheets URL. */
      spreadsheet: SPREADSHEET_REFERENCE_SCHEMA,
    }),
  )
  .output(
    z.object({
      /** Number of cells written across inserts and updates. */
      affectedCells: z.number().int().nonnegative(),

      /** Number of new rows appended. */
      insertedRows: z.number().int().nonnegative(),

      /** A1 range appended for ordinary sheets, when supplied by Google. */
      insertedRange: z.string().optional(),

      /** Existing one-based row numbers updated in place. */
      updatedRowNumbers: z.number().int().positive().array(),
    }),
  )
  .retry({ replaySafety: "unsafe" })
  .handler(async ({ account, input }) => {
    const sheetsApi = getGoogleSheetsApi(account.secret)
    const spreadsheetId = normalizeSpreadsheetId(input.spreadsheet)
    const source = await resolveRowSource(
      sheetsApi,
      spreadsheetId,
      input.source,
    )
    const unknownKeyColumns = input.keyColumns.filter(
      (column) => !source.headers.includes(column),
    )
    if (unknownKeyColumns.length) {
      throw new Error(
        `Unknown Google Sheets key columns: ${unknownKeyColumns.join(", ")}.`,
      )
    }

    const inputRows = Array.isArray(input.rows) ? input.rows : [input.rows]
    inputRows.forEach((row) => {
      normalizeRowValues(row, source.headers)
      const missingKeys = input.keyColumns.filter(
        (column) => row[column] === undefined || row[column] === null,
      )
      if (missingKeys.length) {
        throw new Error(
          `Upsert rows require non-null values for key columns: ${missingKeys.join(", ")}.`,
        )
      }
    })

    // Keep this named so the map's value type stays coupled to normalized rows.
    const existingRows = await readSourceRows(sheetsApi, spreadsheetId, source)
    const existingByKey = new Map<string, (typeof existingRows)[number]>()
    existingRows.forEach((row) => {
      const key = getKey(row.values, input.keyColumns)
      if (existingByKey.has(key)) {
        throw new Error(
          `Google Sheets upsert key ${key} matches multiple existing rows.`,
        )
      }
      existingByKey.set(key, row)
    })
    const inputKeys = inputRows.map((row) => getKey(row, input.keyColumns))
    if (new Set(inputKeys).size !== inputKeys.length) {
      throw new Error("Google Sheets upsert input contains duplicate keys.")
    }

    const updates = inputRows.flatMap((row, index) => {
      const existing = existingByKey.get(inputKeys[index]!)
      return existing ? [{ existing, values: row }] : []
    })
    const inserts = inputRows.filter(
      (_, index) => !existingByKey.has(inputKeys[index]!),
    )
    const updateRequests = updates.flatMap(({ existing, values }) =>
      Object.entries(values).map(([header, value]) => {
        const columnIndex = source.headers.indexOf(header)
        return {
          updateCells: {
            fields: "userEnteredValue",
            range: {
              endColumnIndex: columnIndex + 1,
              endRowIndex: existing.rowNumber,
              sheetId: source.sheetId,
              startColumnIndex: columnIndex,
              startRowIndex: existing.rowNumber - 1,
            },
            rows: [
              {
                values: [toGoogleCellData(value, input.inputMode)],
              },
            ],
          },
        }
      }),
    )
    if (updateRequests.length) {
      await sheetsApi.spreadsheets.batchUpdate({
        requestBody: { requests: updateRequests },
        spreadsheetId,
      })
    }

    let insertedRange: string | undefined
    if (inserts.length && source.tableId) {
      await sheetsApi.spreadsheets.batchUpdate({
        requestBody: {
          requests: [
            {
              appendCells: {
                fields: "userEnteredValue",
                rows: inserts.map((row) => ({
                  values: normalizeRowValues(row, source.headers).map((value) =>
                    toGoogleCellData(value, input.inputMode),
                  ),
                })),
                tableId: source.tableId,
              },
            },
          ],
        },
        spreadsheetId,
      })
    } else if (inserts.length) {
      const { data } = await sheetsApi.spreadsheets.values.append({
        insertDataOption: "INSERT_ROWS",
        range: buildA1Range(source.sheetTitle, {
          endColumn: source.columnCount,
          endRow: source.bodyStartRow - 1,
          startColumn: 1,
          startRow: source.bodyStartRow - 1,
        }),
        requestBody: {
          majorDimension: "ROWS",
          values: inserts.map((row) =>
            normalizeRowValues(row, source.headers).map((value) => value ?? ""),
          ),
        },
        spreadsheetId,
        valueInputOption: input.inputMode === "raw" ? "RAW" : "USER_ENTERED",
      })
      insertedRange = data.updates?.updatedRange ?? undefined
    }

    return {
      affectedCells:
        updateRequests.length + inserts.length * source.columnCount,
      insertedRange,
      insertedRows: inserts.length,
      updatedRowNumbers: updates.map(({ existing }) => existing.rowNumber),
    }
  })

/**
 * Serializes a composite row key without primitive-type collisions.
 *
 * @param row - Header-keyed row values.
 * @param keyColumns - Headers forming the composite key.
 */
function getKey(
  row: Record<string, string | number | boolean | null>,
  keyColumns: string[],
) {
  return JSON.stringify(keyColumns.map((column) => row[column] ?? null))
}