import {
  DEFAULT_GATEWAY_MODEL_ID,
  generationInputSchema,
  generationResultSchema,
  type GenerationResult as BaseGenerationResult,
  type GatewayModelId,
} from "@automate.ax/api-contract/runtime"
import {
  encodableSchema,
  type Encodable,
  type ProducingSchema,
} from "@automate.ax/codec"
import { createAnthropic } from "@ai-sdk/anthropic"
import { createCerebras } from "@ai-sdk/cerebras"
import { createCohere } from "@ai-sdk/cohere"
import { createDeepInfra } from "@ai-sdk/deepinfra"
import { createDeepSeek } from "@ai-sdk/deepseek"
import { createFireworks } from "@ai-sdk/fireworks"
import { createGoogleGenerativeAI } from "@ai-sdk/google"
import { createGroq } from "@ai-sdk/groq"
import { createHuggingFace } from "@ai-sdk/huggingface"
import { createMistral } from "@ai-sdk/mistral"
import { createOpenAI } from "@ai-sdk/openai"
import { createPerplexity } from "@ai-sdk/perplexity"
import { createTogetherAI } from "@ai-sdk/togetherai"
import { createXai } from "@ai-sdk/xai"
import { createOpenRouter } from "@openrouter/ai-sdk-provider"
import type { StandardSchemaV1 } from "@standard-schema/spec"
import { createGateway, generateText, Output, type LanguageModel } from "ai"
import type { DistributedOmit } from "type-fest"
import * as z from "zod"
import { defineAction, type ActionObjectInput } from "../../automation/actions"
import type {
  IntegrationAccountReference,
  ResolvedIntegrationAccount,
} from "../../automation/integrations"
import { serializedIntegrationAccountDefinition } from "../../automation/integrations"
import type { Signal } from "../../automation/signal-protocol"

const AI_ACCOUNT_SERVICE_IDS = [
  "anthropic",
  "cerebras",
  "cohere",
  "deepinfra",
  "deepseek",
  "fireworks",
  "google-genai",
  "groq",
  "huggingface",
  "mistral",
  "openai",
  "openrouter",
  "perplexity",
  "togetherai",
  "vercel-ai-gateway",
  "xai",
] as const

const PLATFORM_AI_CREDENTIAL = Symbol("platform-ai-credential")
const JSON_OBJECT_SCHEMA = z.record(z.string(), z.json())

type AiAccountServiceId = (typeof AI_ACCOUNT_SERVICE_IDS)[number]
type GenerateSchema = ProducingSchema<Encodable>

interface AiModelIdsByService {
  anthropic: Parameters<ReturnType<typeof createAnthropic>>[0]
  cerebras: Parameters<ReturnType<typeof createCerebras>>[0]
  cohere: Parameters<ReturnType<typeof createCohere>>[0]
  deepinfra: Parameters<ReturnType<typeof createDeepInfra>>[0]
  deepseek: Parameters<ReturnType<typeof createDeepSeek>>[0]
  fireworks: Parameters<ReturnType<typeof createFireworks>>[0]
  "google-genai": Parameters<ReturnType<typeof createGoogleGenerativeAI>>[0]
  groq: Parameters<ReturnType<typeof createGroq>>[0]
  huggingface: Parameters<ReturnType<typeof createHuggingFace>>[0]
  mistral: Parameters<ReturnType<typeof createMistral>>[0]
  openai: Parameters<ReturnType<typeof createOpenAI>>[0]
  openrouter: never
  perplexity: Parameters<ReturnType<typeof createPerplexity>>[0]
  togetherai: Parameters<ReturnType<typeof createTogetherAI>>[0]
  "vercel-ai-gateway": GatewayModelId
  xai: Parameters<ReturnType<typeof createXai>>[0]
}

/** Model IDs suggested for one provider, with support for newly released IDs. */
export type GenerateModelId<
  TServiceId extends AiAccountServiceId = AiAccountServiceId,
> = AiModelIdsByService[TServiceId] | (string & {})

const DEFAULT_AI_MODELS = {
  anthropic: "claude-haiku-4-5",
  cerebras: "llama3.1-8b",
  cohere: "command-a-03-2025",
  deepinfra: "meta-llama/Llama-3.3-70B-Instruct-Turbo",
  deepseek: "deepseek-chat",
  fireworks: "accounts/fireworks/models/llama-v3p3-70b-instruct",
  "google-genai": "gemini-2.5-flash",
  groq: "llama-3.1-8b-instant",
  huggingface: "meta-llama/Llama-3.1-8B-Instruct",
  mistral: "mistral-small-latest",
  openai: "gpt-4.1-nano",
  openrouter: "openrouter/auto",
  perplexity: "sonar",
  togetherai: "meta-llama/Llama-3.3-70B-Instruct-Turbo",
  "vercel-ai-gateway": DEFAULT_GATEWAY_MODEL_ID,
  xai: "grok-latest",
} as const satisfies Record<AiAccountServiceId, GenerateModelId>

const AI_MODEL_FACTORIES = {
  anthropic: (apiKey: string, model: GenerateModelId) =>
    createAnthropic({ apiKey })(model),
  cerebras: (apiKey: string, model: GenerateModelId) =>
    createCerebras({ apiKey })(model),
  cohere: (apiKey: string, model: GenerateModelId) =>
    createCohere({ apiKey })(model),
  deepinfra: (apiKey: string, model: GenerateModelId) =>
    createDeepInfra({ apiKey })(model),
  deepseek: (apiKey: string, model: GenerateModelId) =>
    createDeepSeek({ apiKey })(model),
  fireworks: (apiKey: string, model: GenerateModelId) =>
    createFireworks({ apiKey })(model),
  "google-genai": (apiKey: string, model: GenerateModelId) =>
    createGoogleGenerativeAI({ apiKey })(model),
  groq: (apiKey: string, model: GenerateModelId) =>
    createGroq({ apiKey })(model),
  huggingface: (apiKey: string, model: GenerateModelId) =>
    createHuggingFace({ apiKey })(model),
  mistral: (apiKey: string, model: GenerateModelId) =>
    createMistral({ apiKey })(model),
  openai: (apiKey: string, model: GenerateModelId) =>
    createOpenAI({ apiKey })(model),
  openrouter: (apiKey: string, model: GenerateModelId) =>
    createOpenRouter({ apiKey })(model),
  perplexity: (apiKey: string, model: GenerateModelId) =>
    createPerplexity({ apiKey })(model),
  togetherai: (apiKey: string, model: GenerateModelId) =>
    createTogetherAI({ apiKey })(model),
  "vercel-ai-gateway": (apiKey: string, model: GenerateModelId) =>
    createGateway({ apiKey })(model),
  xai: (apiKey: string, model: GenerateModelId) => createXai({ apiKey })(model),
} satisfies Record<
  AiAccountServiceId,
  (apiKey: string, model: GenerateModelId) => LanguageModel
>

/** Serializable AI generation result with a schema-dependent output. */
export type GenerateResult<TOutput extends Encodable = string> = Omit<
  BaseGenerationResult,
  "output"
> & {
  /** Generated text or structured value. */
  output: TOutput
}

/** Inputs accepted by the built-in AI generation action. */
export type GenerateOptions<
  TSchema extends GenerateSchema | undefined = undefined,
  TServiceId extends AiAccountServiceId = AiAccountServiceId,
> = DistributedOmit<
  ActionObjectInput<typeof generationInputSchema>,
  "model"
> & {
  /** User-managed provider account. Omit to use the platform AI Gateway. */
  account?: IntegrationAccountReference<TServiceId>

  /** Provider model ID. Defaults according to the selected account. */
  model?:
    | GenerateModelId<NoInfer<TServiceId>>
    | Exclude<
        ActionObjectInput<typeof generationInputSchema>["model"],
        GatewayModelId | undefined
      >

  /** Standard Schema used to constrain and type structured output. */
  schema?: TSchema

  /** Additional model guidance describing the structured output. */
  schemaDescription?: string

  /** Provider-facing name for the structured output. */
  schemaName?: string
}

type PlatformGenerateOptions<
  TSchema extends GenerateSchema | undefined = undefined,
> = DistributedOmit<
  GenerateOptions<TSchema, "vercel-ai-gateway">,
  "account"
> & {
  account?: never
}

/**
 * Generates text or schema-validated structured data.
 *
 * Omit `account` to use the platform AI Gateway. Explicit provider accounts use
 * their own credentials and accept provider-native model IDs.
 *
 * @param options - Prompt, model, generation settings, and optional schema.
 */
export function generate<TSchema extends GenerateSchema>(
  options: PlatformGenerateOptions<TSchema> & { schema: TSchema },
): Signal<GenerateResult<StandardSchemaV1.InferOutput<TSchema>>>

/**
 * Generates structured data through an explicit provider account.
 *
 * @param options - Provider account, prompt, model, and output schema.
 */
export function generate<
  TSchema extends GenerateSchema,
  const TServiceId extends AiAccountServiceId,
>(
  options: GenerateOptions<TSchema, TServiceId> & {
    account: IntegrationAccountReference<TServiceId>
    schema: TSchema
  },
): Signal<GenerateResult<StandardSchemaV1.InferOutput<TSchema>>>

/**
 * Generates structured data when the account may vary at authoring time.
 *
 * @param options - Account, prompt, model, and output schema.
 */
export function generate<TSchema extends GenerateSchema>(
  options: GenerateOptions<TSchema> & { schema: TSchema },
): Signal<GenerateResult<StandardSchemaV1.InferOutput<TSchema>>>

/**
 * Generates plain text through the platform gateway or an explicit provider.
 *
 * @param options - Prompt, model, and generation settings.
 */
export function generate(
  options: PlatformGenerateOptions,
): Signal<GenerateResult>

/**
 * Generates text through an explicit provider account.
 *
 * @param options - Provider account, prompt, and model settings.
 */
export function generate<const TServiceId extends AiAccountServiceId>(
  options: GenerateOptions<undefined, TServiceId> & {
    account: IntegrationAccountReference<TServiceId>
  },
): Signal<GenerateResult>

/**
 * Generates text when the account may vary at authoring time.
 *
 * @param options - Account, prompt, and model settings.
 */
export function generate(options: GenerateOptions): Signal<GenerateResult>

export function generate(
  options: GenerateOptions<GenerateSchema | undefined>,
): unknown {
  const { account, model, schema, schemaDescription, schemaName, ...settings } =
    options
  const runtimeOutput = schema
    ? createRuntimeOutput(schema, schemaName, schemaDescription)
    : undefined
  const action = defineAction("Generate")
    .account(AI_ACCOUNT_SERVICE_IDS, { default: PLATFORM_AI_CREDENTIAL })
    .input(generationInputSchema)
    .output(
      generationResultSchema.extend({
        output: z.custom<Encodable>(),
      }),
    )
    .retry({ replaySafety: "unsafe" })
    .handler(async ({ account: resolvedAccount, input, runtime }) => {
      const output = await runtimeOutput
      if (typeof resolvedAccount === "symbol") {
        const result = await runtime.generate({
          ...input,
          ...(output && { output }),
        })

        return {
          ...result,
          output: schema
            ? await validateGeneratedOutput(schema, result.output)
            : z.string().parse(result.output),
        }
      }

      const { messages, prompt, ...settings } = input
      return normalizeGenerationResult(
        await generateText({
          ...settings,
          ...(prompt === undefined ? { messages } : { prompt }),
          model: createProviderModel(resolvedAccount, input.model),
          output: schema
            ? Output.object({
                description: schemaDescription,
                name: schemaName,
                schema,
              })
            : Output.text(),
        }),
      )
    })

  return action(
    {
      ...settings,
      model:
        model ??
        (account
          ? DEFAULT_AI_MODELS[
              account[serializedIntegrationAccountDefinition].serviceId
            ]
          : DEFAULT_GATEWAY_MODEL_ID),
      // Make structured-output changes invalidate completed action replays.
      ...(runtimeOutput && { __structuredOutput: runtimeOutput }),
    },
    account && { account },
  )
}

/**
 * Resolves one user-managed provider model from its API-key account.
 *
 * @param account - Resolved provider account.
 * @param model - Provider-native model ID.
 */
function createProviderModel(
  account: ResolvedIntegrationAccount<AiAccountServiceId>,
  model: GenerateModelId,
): LanguageModel {
  const { apiKey } = z
    .object({ apiKey: z.string().min(1) })
    .parse(account.secret)

  return AI_MODEL_FACTORIES[account.serviceId](apiKey, model)
}

/**
 * Converts a Standard Schema to the JSON Schema accepted by runtime transport.
 *
 * @param schema - Caller-provided Standard Schema.
 * @param name - Optional provider-facing output name.
 * @param description - Optional structured-output guidance.
 */
async function createRuntimeOutput(
  schema: GenerateSchema,
  name: string | undefined,
  description: string | undefined,
) {
  const responseFormat = await Output.object({ schema }).responseFormat
  if (responseFormat?.type !== "json" || responseFormat.schema === undefined) {
    throw new Error("Structured generation did not produce a JSON schema.")
  }

  return {
    description,
    name,
    schema: JSON_OBJECT_SCHEMA.parse(responseFormat.schema),
  }
}

/**
 * Narrows the platform response through the caller's original schema.
 *
 * @param schema - Caller-provided Standard Schema.
 * @param output - Generated value returned by the platform.
 */
async function validateGeneratedOutput<TSchema extends GenerateSchema>(
  schema: TSchema,
  output: unknown,
): Promise<StandardSchemaV1.InferOutput<TSchema>> {
  const result = await schema["~standard"].validate(output)
  if (result.issues) {
    throw new Error(
      `Generated output validation failed: ${result.issues.map(({ message }) => message).join("; ")}`,
    )
  }
  return result.value
}

/**
 * Selects the serializable AI SDK result fields exposed by the action.
 *
 * @param result - Completed AI SDK generation result.
 */
function normalizeGenerationResult<TOutput extends Encodable>(
  result: Omit<Awaited<ReturnType<typeof generateText>>, "output"> & {
    output: TOutput
  },
): GenerateResult<TOutput> {
  return {
    finishReason: result.finishReason,
    output: result.output,
    providerMetadata: encodableSchema.parse(result.finalStep.providerMetadata),
    rawFinishReason: result.rawFinishReason,
    reasoningText: result.finalStep.reasoningText,
    response: {
      headers: result.finalStep.response.headers,
      id: result.finalStep.response.id,
      modelId: result.finalStep.response.modelId,
      timestamp: result.finalStep.response.timestamp.toISOString(),
    },
    text: result.text,
    usage: {
      ...result.usage,
      raw: encodableSchema.parse(result.usage.raw),
    },
    warnings: result.warnings,
  }
}