import { encodableSchema } from "@automate.ax/codec"
import {
  AXIOM_APL_REQUEST_SCHEMA,
  AXIOM_DATASET_TAG_VALUE_DISCOVERY_SCHEMA,
  AXIOM_EDGE_URL_SCHEMA,
  AXIOM_ID_SCHEMA,
  AXIOM_INGEST_STATUS_SCHEMA,
  AXIOM_METRIC_DISCOVERY_SCHEMA,
  AXIOM_METRIC_TAG_DISCOVERY_SCHEMA,
  AXIOM_METRIC_TAG_VALUE_DISCOVERY_SCHEMA,
  AXIOM_METRICS_INFO_SCHEMA,
  AXIOM_METRICS_RESULT_SCHEMA,
  AXIOM_QUERY_RESULT_SCHEMA,
  AXIOM_TIME_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 EDGE_FIELDS = { edgeUrl: AXIOM_EDGE_URL_SCHEMA }
const EVENT_INPUT_SCHEMA = z
  .record(z.string(), z.unknown())
  .transform((event) => z.record(z.string(), encodableSchema).parse(event))
const METRIC_TIME_FIELDS = {
  endTime: AXIOM_TIME_SCHEMA.prefault("now"),
  startTime: AXIOM_TIME_SCHEMA.prefault("now-30m"),
}
const METRIC_SEARCH_RESPONSE_SCHEMA = z.record(z.string(), z.json())

/**
 * Resolves the RFC3339 window required by Axiom metric-discovery endpoints.
 *
 * @param input - Optional fixed discovery bounds.
 * @param input.endTime - Fixed discovery-window end.
 * @param input.startTime - Fixed discovery-window start.
 */
function metricDiscoveryQuery(input: { endTime?: string; startTime?: string }) {
  const now = new Date()
  return {
    end: input.endTime ?? now.toISOString(),
    start:
      input.startTime ??
      new Date(now.getTime() - 30 * 60 * 1_000).toISOString(),
  }
}

/** Ingests a batch of structured events into one Axiom dataset. */
export const ingestAxiomEvents = defineAction("Ingest Axiom events")
  .describe("Ingests structured JSON events into one Axiom dataset.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(
    z.object({
      ...EDGE_FIELDS,
      datasetName: AXIOM_ID_SCHEMA,
      events: EVENT_INPUT_SCHEMA.array().min(1).max(10_000),
    }),
  )
  .output(AXIOM_INGEST_STATUS_SCHEMA)
  .retry({ replaySafety: "unsafe" })
  .handler(({ account, input }) =>
    getAxiomApi(account, input.edgeUrl).request(
      `ingest/${axiomPath(input.datasetName)}`,
      {
        body: input.events,
        method: "POST",
        responseSchema: AXIOM_INGEST_STATUS_SCHEMA,
        version: "v1",
      },
    ),
  )

/** Runs an APL query against event, log, or trace datasets. */
export const queryAxiomData = defineAction("Query Axiom data")
  .describe("Runs an APL query over a bounded time range.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(AXIOM_APL_REQUEST_SCHEMA.extend(EDGE_FIELDS))
  .output(AXIOM_QUERY_RESULT_SCHEMA)
  .retry({ replaySafety: "safe" })
  .handler(({ account, input: { edgeUrl, ...body } }) =>
    getAxiomApi(account, edgeUrl).request("query/_apl", {
      body,
      method: "POST",
      query: { format: "tabular" },
      responseSchema: AXIOM_QUERY_RESULT_SCHEMA,
      version: "v1",
    }),
  )

/** Runs an MPL query against an Axiom metrics dataset. */
export const queryAxiomMetrics = defineAction("Query Axiom metrics")
  .describe("Runs an MPL query over a bounded metrics time range.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(
    z.object({
      ...EDGE_FIELDS,
      ...METRIC_TIME_FIELDS,
      mpl: z.string().trim().min(1),
    }),
  )
  .output(AXIOM_METRICS_RESULT_SCHEMA)
  .retry({ replaySafety: "safe" })
  .handler(({ account, input: { edgeUrl, ...body } }) =>
    getAxiomApi(account, edgeUrl).request("query/_mpl", {
      body,
      headers: { Accept: "application/vnd.metrics-v2+json" },
      method: "POST",
      responseSchema: AXIOM_METRICS_RESULT_SCHEMA,
      version: "v1",
    }),
  )

/** Lists metric names and their type, temporality, and unit. */
export const listAxiomMetrics = defineAction("List Axiom metrics")
  .describe("Lists metric definitions observed in a metrics dataset.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(AXIOM_METRIC_DISCOVERY_SCHEMA)
  .output(AXIOM_METRICS_INFO_SCHEMA)
  .retry({ replaySafety: "safe" })
  .handler(({ account, input }) =>
    getAxiomApi(account, input.edgeUrl).request(
      `query/metrics/info/datasets/${axiomPath(input.datasetName)}/metrics`,
      {
        headers: { Accept: "application/vnd.metrics-info.v2+json" },
        query: metricDiscoveryQuery(input),
        responseSchema: AXIOM_METRICS_INFO_SCHEMA,
        version: "v1",
      },
    ),
  )

/** Finds metrics associated with one exact observed tag value. */
export const searchAxiomMetrics = defineAction("Search Axiom metrics")
  .describe("Finds metrics associated with one exact tag value.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(
    AXIOM_METRIC_DISCOVERY_SCHEMA.extend({
      value: z.string().trim().min(1),
    }),
  )
  .output(AXIOM_ID_SCHEMA.array())
  .retry({ replaySafety: "safe" })
  .handler(async ({ account, input }) =>
    Object.keys(
      await getAxiomApi(account, input.edgeUrl).request(
        `query/metrics/info/datasets/${axiomPath(input.datasetName)}/metrics`,
        {
          body: { value: input.value },
          method: "POST",
          query: metricDiscoveryQuery(input),
          responseSchema: METRIC_SEARCH_RESPONSE_SCHEMA,
          version: "v1",
        },
      ),
    ),
  )

/** Lists tag names observed in one metrics dataset. */
export const listAxiomMetricTags = defineAction("List Axiom metric tags")
  .describe("Lists tag names observed in one metrics dataset.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(AXIOM_METRIC_DISCOVERY_SCHEMA)
  .output(AXIOM_ID_SCHEMA.array())
  .retry({ replaySafety: "safe" })
  .handler(({ account, input }) =>
    getAxiomApi(account, input.edgeUrl).request(
      `query/metrics/info/datasets/${axiomPath(input.datasetName)}/tags`,
      {
        headers: { Accept: "application/vnd.metrics-info.v2+json" },
        query: metricDiscoveryQuery(input),
        responseSchema: AXIOM_ID_SCHEMA.array(),
        version: "v1",
      },
    ),
  )

/** Lists observed values for one metrics tag. */
export const listAxiomMetricTagValues = defineAction(
  "List Axiom metric tag values",
)
  .describe("Lists observed values for one metrics tag.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(AXIOM_DATASET_TAG_VALUE_DISCOVERY_SCHEMA)
  .output(z.string().array())
  .retry({ replaySafety: "safe" })
  .handler(({ account, input }) =>
    getAxiomApi(account, input.edgeUrl).request(
      `query/metrics/info/datasets/${axiomPath(input.datasetName)}/tags/${axiomPath(input.tag)}/values`,
      {
        headers: { Accept: "application/vnd.metrics-info.v2+json" },
        query: metricDiscoveryQuery(input),
        responseSchema: z.string().array(),
        version: "v1",
      },
    ),
  )

/** Lists tag names observed for one metric. */
export const listAxiomMetricTagsForMetric = defineAction(
  "List Axiom metric tags for metric",
)
  .describe("Lists tag names observed for one metric in a metrics dataset.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(AXIOM_METRIC_TAG_DISCOVERY_SCHEMA)
  .output(AXIOM_ID_SCHEMA.array())
  .retry({ replaySafety: "safe" })
  .handler(({ account, input }) =>
    getAxiomApi(account, input.edgeUrl).request(
      `query/metrics/info/datasets/${axiomPath(input.datasetName)}/metrics/${axiomPath(input.metric)}/tags`,
      {
        headers: { Accept: "application/vnd.metrics-info.v2+json" },
        query: metricDiscoveryQuery(input),
        responseSchema: AXIOM_ID_SCHEMA.array(),
        version: "v1",
      },
    ),
  )

/** Lists observed values for one tag on one metric. */
export const listAxiomMetricTagValuesForMetric = defineAction(
  "List Axiom metric tag values for metric",
)
  .describe("Lists observed values for one tag on one metric.")
  .account(AXIOM_ACCOUNT, AXIOM_API_TOKEN_ACCOUNT)
  .input(AXIOM_METRIC_TAG_VALUE_DISCOVERY_SCHEMA)
  .output(z.string().array())
  .retry({ replaySafety: "safe" })
  .handler(({ account, input }) =>
    getAxiomApi(account, input.edgeUrl).request(
      `query/metrics/info/datasets/${axiomPath(input.datasetName)}/metrics/${axiomPath(input.metric)}/tags/${axiomPath(input.tag)}/values`,
      {
        headers: { Accept: "application/vnd.metrics-info.v2+json" },
        query: metricDiscoveryQuery(input),
        responseSchema: z.string().array(),
        version: "v1",
      },
    ),
  )