import { automationInvokedTriggerDefinition } from "@automate.ax/catalog/triggers/core-invocation"
import { z } from "zod"
import { defineAction } from "../../automation/actions"
import {
  DEFAULT_ACCOUNT_BINDING,
  defineAccount,
  serializedIntegrationAccountDefinition,
  type IntegrationAccountReference,
} from "../../automation/integrations"
import {
  getAutomationPlanningContext,
  getCurrentHookScopePath,
} from "../../automation/runtime"
import {
  branch,
  correlate,
  filter,
  getCurrentContextPrerequisite,
  getCurrentSignalPrerequisites,
  group,
  scope,
  withContextPrerequisite,
  withPrerequisites,
  withoutSignalPrerequisites,
} from "../../automation/signal-operators"
import {
  FailedSignalError,
  transform,
  type DurationString,
  type Signal,
  type SignalOutcome,
} from "../../automation/signal-protocol"
import { createSubscription } from "../../automation/subscription"

const CHATGPT_API_BASE_URL = "https://api.chatgpt.com/v1/workspace_agents/"
const CHATGPT_API_ORIGIN = new URL(CHATGPT_API_BASE_URL).origin
const declareChatGPTAccount = defineAccount("chatgpt")
const CHATGPT_SECRET_SCHEMA = z.object({ apiKey: z.string().min(1) })
const CHATGPT_API_TRIGGER_ID_SCHEMA = z
  .string()
  .regex(/^agtch_.+$/u, "Invalid Workspace Agent API trigger ID")
const CHATGPT_RUN_ID_SCHEMA = z
  .string()
  .regex(/^apirun_.+$/u, "Invalid Workspace Agent run ID")
const CHATGPT_MAX_POLLS_SCHEMA = z.number().int().min(1).max(120).default(60)
const CHATGPT_REQUEST_TIMEOUT_MS = 50_000
const CHATGPT_POLL_ENTRYPOINT_PREFIX = "__$chatgpt-workspace-agent-poll:"
const CHATGPT_POLL_INTERVAL_SCHEMA = z.union([
  z.string().min(1),
  z.number().nonnegative(),
])
const CHATGPT_POLL_DELIVERY_SCHEMA = z.object({
  accountBinding: z.string().min(1),
  apiTriggerId: CHATGPT_API_TRIGGER_ID_SCHEMA,
  maxPolls: CHATGPT_MAX_POLLS_SCHEMA.removeDefault(),
  poll: z.number().int().min(1).max(120),
  pollInterval: CHATGPT_POLL_INTERVAL_SCHEMA,
  runId: CHATGPT_RUN_ID_SCHEMA,
})

const scheduleChatGPTWorkspaceAgentPoll = defineAction(
  "Schedule ChatGPT workspace agent poll",
)
  .input(
    CHATGPT_POLL_DELIVERY_SCHEMA.extend({
      entrypoint: z.string().startsWith(CHATGPT_POLL_ENTRYPOINT_PREFIX),
      runIn: CHATGPT_POLL_INTERVAL_SCHEMA,
    }),
  )
  .output(z.void())
  .retry({ replaySafety: "safe" })
  .handler(async ({ input, runtime }) => {
    const { entrypoint, runIn, ...payload } = input
    await runtime.invokeAutomation({
      automation: runtime.automationId,
      entrypoint,
      payload,
      runIn,
    })
  })

/** States returned while a Workspace Agent trigger run is processed. */
export const CHATGPT_WORKSPACE_AGENT_RUN_STATUS_SCHEMA = z.enum([
  "queued",
  "in_progress",
  "suspended",
  "completed",
  "failed",
])

/** Failure categories returned for terminal Workspace Agent runs. */
export const CHATGPT_WORKSPACE_AGENT_RUN_ERROR_SCHEMA = z.object({
  /** Provider failure category. */
  code: z.enum(["dispatch_failed", "run_failed"]),
})

const CHATGPT_TRIGGER_RESPONSE_SCHEMA = z.object({
  agent_trigger_run_id: CHATGPT_RUN_ID_SCHEMA,
  conversation_url: z.url(),
})
const CHATGPT_RUN_RESPONSE_SCHEMA = z.object({
  agent_id: z.string().min(1),
  api_trigger_id: CHATGPT_API_TRIGGER_ID_SCHEMA,
  conversation_url: z.url(),
  created_at: z.number().int().nonnegative(),
  error: CHATGPT_WORKSPACE_AGENT_RUN_ERROR_SCHEMA.nullable(),
  id: CHATGPT_RUN_ID_SCHEMA,
  object: z.literal("workspace_agent.trigger_run"),
  status: CHATGPT_WORKSPACE_AGENT_RUN_STATUS_SCHEMA,
})
const CHATGPT_ERROR_SCHEMA = z.looseObject({
  error: z
    .union([z.string(), z.looseObject({ message: z.string().optional() })])
    .optional(),
  message: z.string().optional(),
})

/** Result returned after a Workspace Agent trigger is durably accepted. */
export const CHATGPT_WORKSPACE_AGENT_TRIGGER_RESULT_SCHEMA = z.object({
  /** ChatGPT conversation opened for the trigger. */
  conversationUrl: z.url(),

  /** Run identifier used to poll execution status. */
  runId: CHATGPT_RUN_ID_SCHEMA,
})

/** Current state and identity of one Workspace Agent trigger run. */
export const CHATGPT_WORKSPACE_AGENT_RUN_SCHEMA = z.object({
  /** Published Workspace Agent identity. */
  agentId: z.string().min(1),

  /** API trigger channel identity. */
  apiTriggerId: CHATGPT_API_TRIGGER_ID_SCHEMA,

  /** ChatGPT conversation associated with the run. */
  conversationUrl: z.url(),

  /** Time at which ChatGPT accepted the run. */
  createdAt: z.date(),

  /** Failure details for a failed run, otherwise `null`. */
  error: CHATGPT_WORKSPACE_AGENT_RUN_ERROR_SCHEMA.nullable(),

  /** Trigger run identity. */
  id: CHATGPT_RUN_ID_SCHEMA,

  /** Provider object discriminator. */
  object: z.literal("workspace_agent.trigger_run"),

  /** Current execution state. */
  status: CHATGPT_WORKSPACE_AGENT_RUN_STATUS_SCHEMA,
})
type ChatGPTWorkspaceAgentRun = z.output<
  typeof CHATGPT_WORKSPACE_AGENT_RUN_SCHEMA
>

/** Poll states that must correlate to the accepted submission. */
type ChatGPTWorkspaceAgentPollOutcome = Exclude<
  SignalOutcome<ChatGPTWorkspaceAgentRun>,
  { status: "closed" }
>

/** Tagged scheduler failure selected for correlation propagation. */
type ChatGPTWorkspaceAgentPollSchedulingFailure = Extract<
  SignalOutcome<void>,
  { status: "failed" }
>

/**
 * Triggers a published ChatGPT Workspace Agent through its API channel.
 *
 * The action always requests beta run tracking so its result can be passed to
 * {@link getChatGPTWorkspaceAgentRun}.
 */
export const triggerChatGPTWorkspaceAgent = defineAction(
  "Trigger ChatGPT workspace agent",
)
  .describe("Queues a published ChatGPT Workspace Agent run.")
  .account("chatgpt")
  .input(
    z.object({
      /** Stable `agtch_` identifier for the agent's published API channel. */
      apiTriggerId: CHATGPT_API_TRIGGER_ID_SCHEMA,

      /** Stable key that continues an existing agent conversation. */
      conversationKey: z.string().min(1).optional(),

      /** Message passed to the workspace agent. */
      input: z.string().min(1),

      /** Key reused only when retrying the same trigger event. */
      idempotencyKey: z.string().min(1).optional(),
    }),
  )
  .output(CHATGPT_WORKSPACE_AGENT_TRIGGER_RESULT_SCHEMA)
  .retry({
    replaySafety: ({ idempotencyKey }) => (idempotencyKey ? "safe" : "unsafe"),
  })
  .handler(async ({ account, input }) => {
    const response = CHATGPT_TRIGGER_RESPONSE_SCHEMA.parse(
      await (
        await requestChatGPTApi(
          account.secret,
          `${encodeURIComponent(input.apiTriggerId)}/trigger`,
          {
            body: JSON.stringify({
              input: input.input,
              ...(input.conversationKey && {
                conversation_key: input.conversationKey,
              }),
            }),
            headers: {
              "OpenAI-Beta": "workspace_agent_runs=v1",
              ...(input.idempotencyKey && {
                "Idempotency-Key": input.idempotencyKey,
              }),
            },
            method: "POST",
          },
        )
      ).json(),
    )

    return {
      conversationUrl: response.conversation_url,
      runId: response.agent_trigger_run_id,
    }
  })

/** Gets the current status of a ChatGPT Workspace Agent trigger run. */
export const getChatGPTWorkspaceAgentRun = defineAction(
  "Get ChatGPT workspace agent run",
)
  .describe("Gets the latest state of one Workspace Agent trigger run.")
  .account("chatgpt")
  .input(
    z.object({
      /** Stable `agtch_` identifier used to trigger the run. */
      apiTriggerId: CHATGPT_API_TRIGGER_ID_SCHEMA,

      /** `apirun_` identifier returned when the run was triggered. */
      runId: CHATGPT_RUN_ID_SCHEMA,
    }),
  )
  .output(CHATGPT_WORKSPACE_AGENT_RUN_SCHEMA)
  .retry({ replaySafety: "safe" })
  .handler(async ({ account, input }) => {
    const response = CHATGPT_RUN_RESPONSE_SCHEMA.parse(
      await (
        await requestChatGPTApi(
          account.secret,
          `${encodeURIComponent(input.apiTriggerId)}/runs/${encodeURIComponent(input.runId)}`,
        )
      ).json(),
    )

    return {
      agentId: response.agent_id,
      apiTriggerId: response.api_trigger_id,
      conversationUrl: response.conversation_url,
      createdAt: new Date(response.created_at * 1_000),
      error: response.error,
      id: response.id,
      object: response.object,
      status: response.status,
    }
  })

/** Input for triggering and durably waiting on a Workspace Agent run. */
export type RunChatGPTWorkspaceAgentInput = Extract<
  Parameters<typeof triggerChatGPTWorkspaceAgent>[0],
  { apiTriggerId: unknown }
> & {
  /** Maximum number of status checks before the wait fails. Defaults to 60. */
  maxPolls?: number

  /** Durable delay between status checks. Defaults to `"30s"`. */
  pollInterval?: DurationString | number
}

/** Account selection for composed Workspace Agent helpers. */
export interface ChatGPTWorkspaceAgentAccountOptions {
  /** Static project account binding reused by every polling root. */
  account?: string | IntegrationAccountReference<"chatgpt">
}

/**
 * Triggers a Workspace Agent and continues after it reaches a terminal state.
 *
 * No provider invocation remains open between checks. Each interval uses a
 * durable delivery root and a correlated child context. The returned signal
 * emits completed or failed run metadata, not the agent's response.
 *
 * @param input - Trigger fields and durable polling configuration.
 * @param options - Optional ChatGPT account selection.
 */
export function runChatGPTWorkspaceAgent(
  input: RunChatGPTWorkspaceAgentInput,
  options?: ChatGPTWorkspaceAgentAccountOptions,
): Signal<ChatGPTWorkspaceAgentRun> {
  const prerequisites = getCurrentSignalPrerequisites()
  const contextPrerequisite = getCurrentContextPrerequisite()
  return group(
    { name: "Run ChatGPT workspace agent", presentation: "hidden" },
    () =>
      withoutSignalPrerequisites(() => {
        const { maxPolls, pollInterval, ...submission } = input
        const accountBinding = getStaticChatGPTAccountBinding(options)
        declareChatGPTPollingAccount(accountBinding)
        const trigger = () =>
          triggerChatGPTWorkspaceAgent(submission, { account: accountBinding })
        const triggerWithContext = () =>
          contextPrerequisite
            ? withContextPrerequisite(contextPrerequisite, trigger)
            : trigger()
        const accepted = prerequisites
          ? scope(() => withPrerequisites(prerequisites, triggerWithContext))
          : triggerWithContext()
        const entrypoint = `${CHATGPT_POLL_ENTRYPOINT_PREFIX}${getCurrentHookScopePath().join(".")}`
        const poll = createSubscription<
          z.output<typeof CHATGPT_POLL_DELIVERY_SCHEMA>
        >(automationInvokedTriggerDefinition, { entrypoint }, undefined, {
          inferActionBoundary: false,
        })
        const run = getChatGPTWorkspaceAgentRun(
          { apiTriggerId: poll.apiTriggerId, runId: poll.runId },
          { account: poll.accountBinding },
        )
        const canPollAgain = transform(
          [poll.poll, poll.maxPolls],
          (current, maximum) => current < maximum,
        )
        // Named because every terminal and operational failure joins correlation.
        const pollOutcome = branch(
          run.status.transform(
            (status) => status !== "completed" && status !== "failed",
          ),
          () =>
            branch(
              canPollAgain,
              () => {
                // Keep the scheduler signal named while observing its tagged outcome.
                const scheduled = scheduleChatGPTWorkspaceAgentPoll({
                  accountBinding: poll.accountBinding,
                  apiTriggerId: poll.apiTriggerId,
                  entrypoint,
                  maxPolls: poll.maxPolls,
                  poll: poll.poll.transform((value) => value + 1),
                  pollInterval: poll.pollInterval,
                  runId: poll.runId,
                  runIn: poll.pollInterval,
                })
                return filter(
                  scheduled.outcome(),
                  (
                    result,
                  ): result is ChatGPTWorkspaceAgentPollSchedulingFailure =>
                    result.status === "failed",
                ).transform((result) => {
                  throw new FailedSignalError(result.failure)
                })
              },
              () =>
                transform(
                  [run, poll.poll, poll.maxPolls, poll.pollInterval],
                  (current, count, maximum, interval) => {
                    throw new Error(
                      `ChatGPT Workspace Agent run ${current.id} remained ${current.status} after ${count} of ${maximum} checks at ${interval} intervals.`,
                    )
                  },
                ),
            ),
          () => run,
        )
        const offered = transform(
          [
            filter(
              pollOutcome.outcome(),
              (result): result is ChatGPTWorkspaceAgentPollOutcome =>
                result.status !== "closed",
            ),
            poll.runId,
          ],
          (result, runId) => ({ result, runId }),
        )
        return transform(
          [
            correlate([
              transform(
                [
                  accepted,
                  scheduleChatGPTWorkspaceAgentPoll({
                    accountBinding,
                    apiTriggerId: submission.apiTriggerId,
                    entrypoint,
                    maxPolls: CHATGPT_MAX_POLLS_SCHEMA.parse(maxPolls),
                    poll: 1,
                    pollInterval: pollInterval ?? "30s",
                    runId: accepted.runId,
                    runIn: 0,
                  }),
                ],
                (acceptedRun) => acceptedRun,
              ).keyBy(({ runId }) => runId),
              offered.keyBy(({ runId }) => runId),
            ]),
            offered.result,
          ],
          (_match, result) => {
            if (result.status === "failed") {
              throw new FailedSignalError(result.failure)
            }
            return result.value
          },
        )
      }),
  )
}

/**
 * Declares the one binding a dynamic polling delivery may select.
 *
 * @param binding - Static binding persisted in every polling payload.
 */
function declareChatGPTPollingAccount(binding: string) {
  const planning = getAutomationPlanningContext()
  if (
    !planning ||
    planning.accountDeclarations.some(
      (account) =>
        account.serviceId === "chatgpt" && account.binding === binding,
    )
  ) {
    return
  }
  declareChatGPTAccount(binding)
}

/**
 * Resolves the static binding that polling roots use to load credentials.
 *
 * @param options - Optional static ChatGPT account selection.
 */
function getStaticChatGPTAccountBinding(
  options: ChatGPTWorkspaceAgentAccountOptions | undefined,
) {
  return typeof options?.account === "string"
    ? options.account
    : (options?.account?.[serializedIntegrationAccountDefinition].binding ??
        DEFAULT_ACCOUNT_BINDING)
}

/**
 * Sends one authenticated Workspace Agents API request.
 *
 * @param secret - Resolved ChatGPT integration secret.
 * @param path - Relative Workspace Agents API path.
 * @param options - Native fetch options.
 */
async function requestChatGPTApi(
  secret: Record<string, unknown>,
  path: string,
  options: Omit<RequestInit, "signal"> = {},
) {
  const { apiKey } = CHATGPT_SECRET_SCHEMA.parse(secret)
  const url = new URL(path.replace(/^\//, ""), CHATGPT_API_BASE_URL)
  if (url.origin !== CHATGPT_API_ORIGIN) {
    throw new Error("ChatGPT API paths must use the ChatGPT API origin.")
  }

  const headers = new Headers(options.headers)
  headers.set("Authorization", `Bearer ${apiKey}`)
  headers.set("Accept", "application/json")
  if (options.body !== undefined) {
    headers.set("Content-Type", "application/json")
  }

  const response = await fetch(url, {
    ...options,
    headers,
    signal: AbortSignal.timeout(CHATGPT_REQUEST_TIMEOUT_MS),
  })
  if (response.ok) return response

  const error = CHATGPT_ERROR_SCHEMA.safeParse(
    await response.json().catch(() => ({})),
  )
  const message =
    error.success && typeof error.data.error === "object"
      ? error.data.error.message
      : error.success && typeof error.data.error === "string"
        ? error.data.error
        : error.success
          ? error.data.message
          : undefined
  throw new Error(
    message
      ? `ChatGPT API error (${response.status}): ${message}`
      : `ChatGPT API request failed with status ${response.status}.`,
  )
}