import type { Client } from "@notionhq/client"
import * as z from "zod"
import { notionAccount, notionAction, notionOutput } from "../lib"

const AGENT_STATUS_SCHEMA = z.enum(["active", "disabled"])
const AGENT_BATCH_INPUT_SCHEMA = z.strictObject({
  operations: z
    .discriminatedUnion("action", [
      z.strictObject({
        action: z.literal("update_status"),
        agent_id: z.string().min(1),
        fields: z.strictObject({ status: AGENT_STATUS_SCHEMA }),
      }),
      z.strictObject({
        action: z.literal("update_credit_limit"),
        agent_id: z.string().min(1),
        fields: z.strictObject({
          credit_limit: z.number().finite().nullable(),
        }),
      }),
      z.strictObject({
        action: z.literal("delete"),
        agent_id: z.string().min(1),
      }),
    ])
    .array(),
})
const SESSION_ATTACHMENT_SCHEMA = z.strictObject({
  file_upload: z.strictObject({ id: z.string().min(1) }),
  name: z.string().optional(),
  type: z.literal("file_upload").optional(),
})
const SESSION_INPUT_SCHEMA = z.union([
  z.strictObject({
    actions: z
      .strictObject({
        action_id: z.string().min(1),
        option_id: z.enum(["approve", "reject"]),
      })
      .array(),
    metadata: z.record(z.string(), z.string()).optional(),
    session_id: z.string().min(1),
  }),
  z.strictObject({
    agent_id: z.string().min(1).optional(),
    attachments: SESSION_ATTACHMENT_SCHEMA.array().optional(),
    message: z.string(),
    metadata: z.record(z.string(), z.string()).optional(),
    prompt_context: z.string().optional(),
    session_id: z.string().min(1).optional(),
  }),
  z.strictObject({
    continue_from: z.string(),
    session_id: z.string().min(1),
  }),
])

/** Queries Notion agents available to the connected workspace. */
export const queryNotionAgents = notionAction<Client["agents"]["query"]>({
  account: notionAccount.agents,
  allowedKeys: [
    "query",
    "filter",
    "sorts",
    "start_cursor",
    "page_size",
    "verbose",
    "include_deleted",
  ],
  call: async (client, input) => await client.agents.query(input),
  description: "Queries accessible Notion agents with filters and pagination.",
  name: "Query Notion agents",
  replaySafety: "safe",
  output: notionOutput.list("agent"),
})

/** Retrieves a Notion agent. */
export const getNotionAgent = notionAction<Client["agents"]["retrieve"]>({
  account: notionAccount.agents,
  allowedKeys: ["agent_id", "verbose"],
  call: async (client, input) => await client.agents.retrieve(input),
  description: "Retrieves an agent and optionally its verbose configuration.",
  name: "Get Notion agent",
  replaySafety: "safe",
  output: notionOutput.entity("agent"),
  requiredKeys: ["agent_id"],
})

/** Retrieves usage insights for a Notion agent. */
export const getNotionAgentInsights = notionAction<
  Client["agents"]["retrieveInsights"]
>({
  account: notionAccount.agents,
  allowedKeys: ["agent_id", "start_time", "end_time"],
  call: async (client, input) => await client.agents.retrieveInsights(input),
  description:
    "Retrieves agent usage and run insights for an optional time range.",
  name: "Get Notion agent insights",
  replaySafety: "safe",
  output: notionOutput.entity("agent_insights", [
    "id",
    "name",
    "agent_type",
    "status",
    "pause_reason",
    "created_by",
    "total_credits_used",
    "credit_limit",
    "runs_completed",
  ]),
  requiredKeys: ["agent_id"],
})

/** Enables or disables a Notion agent. */
export const updateNotionAgentStatus = notionAction<
  Client["agents"]["updateStatus"]
>({
  account: notionAccount.agents,
  allowedKeys: ["agent_id", "status"],
  call: async (client, input) => await client.agents.updateStatus(input),
  description: "Enables or disables an agent.",
  inputSchema: z.strictObject({
    agent_id: z.string().min(1),
    status: AGENT_STATUS_SCHEMA,
  }),
  name: "Update Notion agent status",
  replaySafety: "unsafe",
  output: notionOutput.fields(
    "agent_id",
    "status",
    "pause_reason",
    "last_edited_time",
  ),
  requiredKeys: ["agent_id", "status"],
})

/** Updates a Notion agent's credit limit. */
export const updateNotionAgentCreditLimit = notionAction<
  Client["agents"]["updateCreditLimit"]
>({
  account: notionAccount.agents,
  allowedKeys: ["agent_id", "credit_limit"],
  call: async (client, input) => await client.agents.updateCreditLimit(input),
  description: "Sets or clears an agent credit limit.",
  inputSchema: z.strictObject({
    agent_id: z.string().min(1),
    credit_limit: z.number().finite().nullable(),
  }),
  name: "Update Notion agent credit limit",
  replaySafety: "unsafe",
  output: notionOutput.fields("agent_id", "credit_limit", "last_edited_time"),
  requiredKeys: ["agent_id", "credit_limit"],
})

/** Deletes a Notion agent. */
export const deleteNotionAgent = notionAction<Client["agents"]["delete"]>({
  account: notionAccount.agents,
  allowedKeys: ["agent_id"],
  call: async (client, input) => await client.agents.delete(input),
  description: "Deletes one agent.",
  name: "Delete Notion agent",
  replaySafety: "unsafe",
  output: notionOutput.fields("agent_id", "status", "deleted_at"),
  requiredKeys: ["agent_id"],
})

/** Applies multiple status, credit, or delete operations to Notion agents. */
export const batchManageNotionAgents = notionAction<Client["agents"]["batch"]>({
  account: notionAccount.agents,
  allowedKeys: ["operations"],
  call: async (client, input) => await client.agents.batch(input),
  description: "Applies multiple status, credit-limit, or delete operations.",
  inputSchema: AGENT_BATCH_INPUT_SCHEMA,
  name: "Batch manage Notion agents",
  replaySafety: "unsafe",
  output: notionOutput.entity("async_task", [
    "id",
    "status_url",
    "created_time",
    "operation",
    "status",
    "poll_after_seconds",
  ]),
  requiredKeys: ["operations"],
})

/** Starts a Notion agent session or adds input to an existing session. */
export const createOrUpdateNotionSession = notionAction<
  Client["sessions"]["update"]
>({
  account: notionAccount.agents,
  allowedKeys: [
    "message",
    "agent_id",
    "session_id",
    "attachments",
    "metadata",
    "prompt_context",
    "actions",
    "continue_from",
  ],
  call: async (client, input) => await client.sessions.update(input),
  description:
    "Starts a session or sends a message, action response, or continuation.",
  inputSchema: SESSION_INPUT_SCHEMA,
  name: "Create or update Notion session",
  replaySafety: "unsafe",
  output: notionOutput.entity("session", [
    "id",
    "agent_id",
    "title",
    "status",
    "created_at",
    "updated_at",
  ]),
})

/** Retrieves a Notion agent session. */
export const getNotionSession = notionAction<Client["sessions"]["retrieve"]>({
  account: notionAccount.agents,
  allowedKeys: ["session_id"],
  call: async (client, input) => await client.sessions.retrieve(input),
  description: "Retrieves one agent session and its current state.",
  name: "Get Notion session",
  replaySafety: "safe",
  output: notionOutput.entity("session", [
    "id",
    "agent_id",
    "title",
    "status",
    "created_at",
    "updated_at",
  ]),
  requiredKeys: ["session_id"],
})

/** Queries Notion agent sessions. */
export const queryNotionSessions = notionAction<Client["sessions"]["query"]>({
  account: notionAccount.agents,
  allowedKeys: ["query", "filter", "sorts", "start_cursor", "page_size"],
  call: async (client, input) => await client.sessions.query(input),
  description: "Queries agent sessions with filters, sorts, and pagination.",
  name: "Query Notion sessions",
  replaySafety: "safe",
  output: notionOutput.list("session"),
})

/** Queries events emitted within a Notion agent session. */
export const queryNotionSessionEvents = notionAction<
  Client["sessions"]["queryEvents"]
>({
  account: notionAccount.agents,
  allowedKeys: ["session_id", "filter", "sorts", "start_cursor", "page_size"],
  call: async (client, input) => await client.sessions.queryEvents(input),
  description: "Queries events for one agent session.",
  name: "Query Notion session events",
  replaySafety: "safe",
  output: notionOutput.list("session_event"),
  requiredKeys: ["session_id"],
})

/** Cancels a running Notion agent session. */
export const cancelNotionSession = notionAction<Client["sessions"]["cancel"]>({
  account: notionAccount.agents,
  allowedKeys: ["session_id", "event_id"],
  call: async (client, input) => await client.sessions.cancel(input),
  description:
    "Cancels a running agent session, optionally at an event boundary.",
  name: "Cancel Notion session",
  replaySafety: "unsafe",
  output: notionOutput.entity("session", [
    "id",
    "agent_id",
    "title",
    "status",
    "created_at",
    "updated_at",
  ]),
  requiredKeys: ["session_id"],
})