import type { Beta } from "@anthropic-ai/sdk/resources/beta/beta"
import { anthropicTriggerDefinitions } from "@automate.ax/catalog/triggers/anthropic"
import type { AnthropicRunEvent } from "@automate.ax/integration-contracts/anthropic"
import * as z from "zod"
import { defineAction } from "../../automation/actions"
import type { IntegrationAccountReference } from "../../automation/integrations"
import {
  correlate,
  getCurrentSignalPrerequisites,
  group,
  merge,
  withoutSignalPrerequisites,
  withPrerequisites,
} from "../../automation/signal-operators"
import { transform } from "../../automation/signal-protocol"
import { createSubscription } from "../../automation/subscription"
import {
  anthropicInputSchema,
  anthropicOutputSchema,
  fromAnthropic,
  getAnthropicApi,
  toAnthropic,
} from "./lib"

type InitialSessionEvent = NonNullable<
  Beta.SessionCreateParams["initial_events"]
>[number]

/** Session parameters with initial work required for a run-to-result helper. */
type ClaudeAgentRunCreateParams = Omit<
  Beta.SessionCreateParams,
  "initial_events"
> & {
  initial_events: [InitialSessionEvent, ...InitialSessionEvent[]]
}

const START_CLAUDE_AGENT_RUN = defineAction("Start Claude agent run")
  .describe("Starts a Claude agent session with required initial work.")
  .account("anthropic")
  .input(
    anthropicInputSchema<ClaudeAgentRunCreateParams>(
      z.looseObject({
        agent: z.union([z.string().min(1), z.object({}).loose()]),
        environmentId: z.string().min(1),
        initialEvents: z
          .array(
            z.looseObject({
              type: z.enum(["user.define_outcome", "user.message"]),
            }),
          )
          .min(1)
          .max(50),
      }),
      [
        "agent",
        "budget",
        "environmentId",
        "initialEvents",
        "metadata",
        "resources",
        "title",
        "vaultIds",
      ],
    ),
  )
  .output(anthropicOutputSchema<Beta.BetaManagedAgentsSession>())
  .retry({ replaySafety: "unsafe" })
  .handler(async ({ account, input }) =>
    fromAnthropic(
      await getAnthropicApi(account.secret).beta.sessions.create(
        toAnthropic<ClaudeAgentRunCreateParams>(input),
      ),
    ),
  )

export type ClaudeAgentRunResult =
  | {
      event: AnthropicRunEvent
      session: AnthropicRunEvent["resource"]
      status: "completed" | "budgetReached" | "retriesExhausted"
    }
  | {
      event: AnthropicRunEvent
      pendingEventIds: string[]
      session: AnthropicRunEvent["resource"]
      status: "requiresAction"
    }
  | {
      event: AnthropicRunEvent
      session: AnthropicRunEvent["resource"]
      status: "terminated"
    }

export interface ClaudeAgentRunOptions {
  /** Static Managed Agents connection shared by the action and callbacks. */
  account?: string | IntegrationAccountReference<"anthropic">
}

/**
 * Starts a Claude session and resumes in a correlated child when the first turn
 * idles or terminates.
 *
 * @param input - Managed Agents session creation input.
 * @param options - Static account selection shared by action and callbacks.
 */
export function runClaudeAgent(
  input: Parameters<typeof START_CLAUDE_AGENT_RUN>[0],
  options?: ClaudeAgentRunOptions,
) {
  const prerequisites = getCurrentSignalPrerequisites()
  return group({ name: "Run Claude agent", presentation: "hidden" }, () =>
    withoutSignalPrerequisites(() => {
      const config = options?.account ? { account: options.account } : {}
      const completion = merge([
        createSubscription<AnthropicRunEvent>(
          anthropicTriggerDefinitions.anthropicSessionIdled,
          config,
          undefined,
          { inferActionBoundary: false },
        ),
        createSubscription<AnthropicRunEvent>(
          anthropicTriggerDefinitions.anthropicSessionTerminated,
          config,
          undefined,
          { inferActionBoundary: false },
        ),
      ])
      const start = () => START_CLAUDE_AGENT_RUN(input, options)
      // Preserve a caller boundary only for the initiating action.
      const created = prerequisites
        ? withPrerequisites(prerequisites, start)
        : start()
      // The named join documents that callbacks remain independent roots.
      const resumed = correlate(
        [
          created.keyBy(({ id }) => id),
          completion.keyBy(({ resourceId }) => resourceId),
        ],
        { occurrenceTtl: "30d" },
      )
      return transform(
        transform([resumed, completion], (_resumed, value) => value),
        classifyClaudeAgentRun,
      )
    }),
  )
}

/**
 * Converts a terminal session event into the durable helper result.
 *
 * @param event - Enriched session webhook event.
 */
function classifyClaudeAgentRun(
  event: AnthropicRunEvent,
): ClaudeAgentRunResult {
  if (event.providerEventType === "session.status_terminated") {
    return { event, session: event.resource, status: "terminated" }
  }
  const reason = event.idleEvent.stopReason
  if (reason.type === "requires_action") {
    return {
      event,
      pendingEventIds: reason.eventIds,
      session: event.resource,
      status: "requiresAction",
    }
  }
  return {
    event,
    session: event.resource,
    status:
      reason.type === "end_turn"
        ? "completed"
        : reason.type === "budget_reached"
          ? "budgetReached"
          : "retriesExhausted",
  }
}