import {
  AXIOM_DATASET_FIELD_SCHEMA,
  AXIOM_DATASET_KIND_SCHEMA,
  AXIOM_DATASET_SCHEMA,
  AXIOM_EMPTY_SCHEMA,
  AXIOM_ID_SCHEMA,
} from "@automate.ax/integration-contracts/axiom"
import * as z from "zod"
import { defineAction } from "../../../automation/actions"
import {
  AXIOM_ACCOUNT,
  AXIOM_API_TOKEN_ACCOUNT,
  axiomPath,
  getAxiomApi,
} from "../lib"

const DATASET_NAME = { datasetName: AXIOM_ID_SCHEMA }
const DATASET_SETTINGS = {
  description: z.string().optional(),
  retentionDays: z.number().int().nonnegative().optional(),
  useRetentionPeriod: z.boolean().optional(),
}

export const listAxiomDatasets = defineAction("List Axiom datasets")
  .describe("Lists datasets visible to the connected token.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object({}))
  .output(AXIOM_DATASET_SCHEMA.array())
  .retry({ replaySafety: "safe" })
  .handler(({ account }) =>
    getAxiomApi(account).request("datasets", {
      responseSchema: AXIOM_DATASET_SCHEMA.array(),
    }),
  )

export const createAxiomDataset = defineAction("Create Axiom dataset")
  .describe("Creates an Axiom dataset with optional retention settings.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(
    z.object({
      ...DATASET_SETTINGS,
      edgeDeployment: z.string().optional(),
      kind: AXIOM_DATASET_KIND_SCHEMA.prefault("axiom:events:v1"),
      name: AXIOM_ID_SCHEMA,
    }),
  )
  .output(AXIOM_DATASET_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request("datasets", {
      body: input,
      method: "POST",
      responseSchema: AXIOM_DATASET_SCHEMA,
    }),
  )

export const getAxiomDataset = defineAction("Get Axiom dataset")
  .describe("Retrieves one Axiom dataset by name or ID.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object(DATASET_NAME))
  .output(AXIOM_DATASET_SCHEMA)
  .retry({ replaySafety: "safe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(`datasets/${axiomPath(input.datasetName)}`, {
      responseSchema: AXIOM_DATASET_SCHEMA,
    }),
  )

export const updateAxiomDataset = defineAction("Update Axiom dataset")
  .describe("Updates one Axiom dataset's description or retention settings.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(
    z
      .object({ ...DATASET_NAME, ...DATASET_SETTINGS })
      .refine(
        ({ description, retentionDays, useRetentionPeriod }) =>
          description !== undefined ||
          retentionDays !== undefined ||
          useRetentionPeriod !== undefined,
        { message: "Provide at least one dataset setting to update." },
      ),
  )
  .output(AXIOM_DATASET_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input: { datasetName, ...body } }) =>
    getAxiomApi(account).request(`datasets/${axiomPath(datasetName)}`, {
      body,
      method: "PUT",
      responseSchema: AXIOM_DATASET_SCHEMA,
    }),
  )

export const deleteAxiomDataset = defineAction("Delete Axiom dataset")
  .describe("Permanently deletes one Axiom dataset and its data.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object(DATASET_NAME))
  .output(AXIOM_EMPTY_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(`datasets/${axiomPath(input.datasetName)}`, {
      method: "DELETE",
      responseSchema: AXIOM_EMPTY_SCHEMA,
    }),
  )

export const trimAxiomDataset = defineAction("Trim Axiom dataset")
  .describe("Deletes dataset events older than a relative duration.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object({ ...DATASET_NAME, maxDuration: z.string().trim().min(1) }))
  .output(AXIOM_EMPTY_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(
      `datasets/${axiomPath(input.datasetName)}/trim`,
      {
        body: { maxDuration: input.maxDuration },
        method: "POST",
        responseSchema: AXIOM_EMPTY_SCHEMA,
      },
    ),
  )

export const vacuumAxiomDataset = defineAction("Vacuum Axiom dataset")
  .describe("Removes fields no longer present in retained dataset data.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object(DATASET_NAME))
  .output(AXIOM_EMPTY_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(
      `datasets/${axiomPath(input.datasetName)}/vacuum`,
      { method: "POST", responseSchema: AXIOM_EMPTY_SCHEMA },
    ),
  )

export const listAxiomDatasetFields = defineAction("List Axiom dataset fields")
  .describe("Lists the field schema for one events, logs, or traces dataset.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object(DATASET_NAME))
  .output(AXIOM_DATASET_FIELD_SCHEMA.array())
  .retry({ replaySafety: "safe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(
      `datasets/${axiomPath(input.datasetName)}/fields`,
      { responseSchema: AXIOM_DATASET_FIELD_SCHEMA.array() },
    ),
  )

export const getAxiomDatasetField = defineAction("Get Axiom dataset field")
  .describe("Retrieves one field definition from an Axiom dataset.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object({ ...DATASET_NAME, fieldName: AXIOM_ID_SCHEMA }))
  .output(AXIOM_DATASET_FIELD_SCHEMA)
  .retry({ replaySafety: "safe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(
      `datasets/${axiomPath(input.datasetName)}/fields/${axiomPath(input.fieldName)}`,
      { responseSchema: AXIOM_DATASET_FIELD_SCHEMA },
    ),
  )

export const updateAxiomDatasetField = defineAction(
  "Update Axiom dataset field",
)
  .describe("Updates field visibility, description, type, or unit.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(
    AXIOM_DATASET_FIELD_SCHEMA.extend({
      datasetName: AXIOM_ID_SCHEMA,
      fieldName: AXIOM_ID_SCHEMA,
      name: AXIOM_ID_SCHEMA.optional(),
    }),
  )
  .output(AXIOM_DATASET_FIELD_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input: { datasetName, fieldName, ...body } }) =>
    getAxiomApi(account).request(
      `datasets/${axiomPath(datasetName)}/fields/${axiomPath(fieldName)}`,
      {
        body: { ...body, name: body.name ?? fieldName },
        method: "PUT",
        responseSchema: AXIOM_DATASET_FIELD_SCHEMA,
      },
    ),
  )

export const listAxiomMapFields = defineAction("List Axiom map fields")
  .describe("Lists fields configured as maps in one Axiom dataset.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object(DATASET_NAME))
  .output(AXIOM_ID_SCHEMA.array())
  .retry({ replaySafety: "safe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(
      `datasets/${axiomPath(input.datasetName)}/mapfields`,
      { responseSchema: AXIOM_ID_SCHEMA.array() },
    ),
  )

export const replaceAxiomMapFields = defineAction("Replace Axiom map fields")
  .describe("Replaces the complete map-field list for one dataset.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object({ ...DATASET_NAME, fieldNames: AXIOM_ID_SCHEMA.array() }))
  .output(AXIOM_ID_SCHEMA.array())
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(
      `datasets/${axiomPath(input.datasetName)}/mapfields`,
      {
        body: input.fieldNames,
        method: "PUT",
        responseSchema: AXIOM_ID_SCHEMA.array(),
      },
    ),
  )

export const createAxiomMapField = defineAction("Create Axiom map field")
  .describe("Marks one dataset field as a map field.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object({ ...DATASET_NAME, fieldName: AXIOM_ID_SCHEMA }))
  .output(z.object({ name: AXIOM_ID_SCHEMA }))
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(
      `datasets/${axiomPath(input.datasetName)}/mapfields`,
      {
        body: { name: input.fieldName },
        method: "POST",
        responseSchema: z.object({ name: AXIOM_ID_SCHEMA }),
      },
    ),
  )

export const deleteAxiomMapField = defineAction("Delete Axiom map field")
  .describe("Removes one field from a dataset's map-field configuration.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(z.object({ ...DATASET_NAME, fieldName: AXIOM_ID_SCHEMA }))
  .output(AXIOM_EMPTY_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input }) =>
    getAxiomApi(account).request(
      `datasets/${axiomPath(input.datasetName)}/mapfields/${axiomPath(input.fieldName)}`,
      { method: "DELETE", responseSchema: AXIOM_EMPTY_SCHEMA },
    ),
  )