import {
  DEFAULT_GATEWAY_EVALUATION_MODEL_ID,
  evaluationInputSchema,
  evaluationResultSchema,
  type EvaluationInput,
  type EvaluationResult as BaseEvaluationResult,
  type GatewayEvaluationModelId,
} from "@automate.ax/api-contract/runtime"
import { createAnthropic } from "@ai-sdk/anthropic"
import { createGoogleGenerativeAI } from "@ai-sdk/google"
import { createOpenAI } from "@ai-sdk/openai"
import {
  createGateway,
  experimental_evaluate as evaluateAi,
  type Experimental_EvaluationModel,
  type Experimental_EvaluationResult,
} from "ai"
import type { DistributedOmit } from "type-fest"
import { 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_EVALUATION_ACCOUNT_SERVICE_IDS = [
  "anthropic",
  "google-genai",
  "openai",
  "vercel-ai-gateway",
] as const

const PLATFORM_AI_CREDENTIAL = Symbol("platform-ai-credential")

type AiEvaluationAccountServiceId =
  (typeof AI_EVALUATION_ACCOUNT_SERVICE_IDS)[number]

interface AiEvaluationModelIdsByService {
  anthropic: Parameters<
    ReturnType<typeof createAnthropic>["evaluationModel"]
  >[0]
  "google-genai": Parameters<
    ReturnType<typeof createGoogleGenerativeAI>["evaluationModel"]
  >[0]
  openai: Parameters<ReturnType<typeof createOpenAI>["evaluationModel"]>[0]
  "vercel-ai-gateway": GatewayEvaluationModelId
}

/** Model IDs suggested for one evaluation provider, including new string IDs. */
export type EvaluateModelId<
  TServiceId extends
    AiEvaluationAccountServiceId = AiEvaluationAccountServiceId,
> = AiEvaluationModelIdsByService[TServiceId] | (string & {})

/** One Choice, Score, or Boolean question evaluated against shared state. */
export type EvaluateQuestion = EvaluationInput["questions"][string]

/** Named questions evaluated together against shared state. */
export type EvaluateQuestions = Record<string, EvaluateQuestion>

const DEFAULT_EVALUATION_MODELS = {
  anthropic: "claude-haiku-4-5-20251001",
  "google-genai": "gemini-3.5-flash-lite",
  openai: "gpt-5.6-luna",
  "vercel-ai-gateway": DEFAULT_GATEWAY_EVALUATION_MODEL_ID,
} as const satisfies Record<AiEvaluationAccountServiceId, EvaluateModelId>

const AI_EVALUATION_MODEL_FACTORIES = {
  anthropic: (apiKey: string, model: EvaluateModelId) =>
    createAnthropic({ apiKey }).evaluationModel(model),
  "google-genai": (apiKey: string, model: EvaluateModelId) =>
    createGoogleGenerativeAI({ apiKey }).evaluationModel(model),
  openai: (apiKey: string, model: EvaluateModelId) =>
    createOpenAI({ apiKey }).evaluationModel(model),
  "vercel-ai-gateway": (apiKey: string, model: EvaluateModelId) =>
    createGateway({ apiKey }).evaluationModel(model),
} satisfies Record<
  AiEvaluationAccountServiceId,
  (apiKey: string, model: EvaluateModelId) => Experimental_EvaluationModel
>

/** Serializable evaluation result with question-dependent answer types. */
export type EvaluateResult<TQuestions extends EvaluateQuestions> = Omit<
  BaseEvaluationResult,
  "answers"
> & {
  /** One typed answer for every named question. */
  answers: Experimental_EvaluationResult<TQuestions>["answers"]
}

/** Inputs accepted by the built-in AI evaluation action. */
export type EvaluateOptions<
  TQuestions extends EvaluateQuestions,
  TServiceId extends
    AiEvaluationAccountServiceId = AiEvaluationAccountServiceId,
> = DistributedOmit<
  ActionObjectInput<typeof evaluationInputSchema>,
  "model" | "questions"
> & {
  /** User-managed provider account. Omit to use the platform AI Gateway. */
  account?: IntegrationAccountReference<TServiceId>

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

  /** Static named questions whose keys and choices type the returned answers. */
  questions: TQuestions
}

/**
 * Evaluates typed questions through the platform AI Gateway.
 *
 * @param options - Shared state, static questions, model, and evaluation
 *   settings.
 */
export function evaluate<const TQuestions extends EvaluateQuestions>(
  options: DistributedOmit<
    EvaluateOptions<TQuestions, "vercel-ai-gateway">,
    "account"
  > & {
    account?: never
  },
): Signal<EvaluateResult<TQuestions>>

/**
 * Evaluates typed questions through an explicit provider account.
 *
 * @param options - Provider account, shared state, questions, and model
 *   settings.
 */
export function evaluate<
  const TQuestions extends EvaluateQuestions,
  const TServiceId extends AiEvaluationAccountServiceId,
>(
  options: EvaluateOptions<TQuestions, TServiceId> & {
    account: IntegrationAccountReference<TServiceId>
  },
): Signal<EvaluateResult<TQuestions>>

/**
 * Evaluates typed questions when the account may vary at authoring time.
 *
 * @param options - Shared state, static questions, account, and model settings.
 */
export function evaluate<const TQuestions extends EvaluateQuestions>(
  options: EvaluateOptions<TQuestions>,
): Signal<EvaluateResult<TQuestions>>

export function evaluate<const TQuestions extends EvaluateQuestions>(
  options: EvaluateOptions<TQuestions>,
): unknown {
  const { account, model, ...settings } = options
  const action = defineAction("Evaluate")
    .account(AI_EVALUATION_ACCOUNT_SERVICE_IDS, {
      default: PLATFORM_AI_CREDENTIAL,
    })
    .input(evaluationInputSchema)
    .output(evaluationResultSchema)
    .retry({ replaySafety: "unsafe" })
    .handler(async ({ account: resolvedAccount, input, runtime }) => {
      if (typeof resolvedAccount === "symbol") {
        return await runtime.evaluate(input)
      }

      return normalizeEvaluationResult(
        await evaluateAi({
          ...input,
          model: createProviderEvaluationModel(resolvedAccount, input.model),
        }),
      )
    })

  return action(
    {
      ...settings,
      model:
        model ??
        (account
          ? DEFAULT_EVALUATION_MODELS[
              account[serializedIntegrationAccountDefinition].serviceId
            ]
          : DEFAULT_GATEWAY_EVALUATION_MODEL_ID),
    },
    account && { account },
  )
}

/**
 * Resolves one user-managed provider evaluation model.
 *
 * @param account - Resolved provider account credentials.
 * @param model - Provider evaluation model identifier.
 */
function createProviderEvaluationModel(
  account: ResolvedIntegrationAccount<AiEvaluationAccountServiceId>,
  model: EvaluateModelId,
): Experimental_EvaluationModel {
  const { apiKey } = z
    .object({ apiKey: z.string().min(1) })
    .parse(account.secret)

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

/**
 * Selects the serializable AI SDK result fields exposed by the action.
 *
 * @param result - Raw AI SDK evaluation result.
 */
function normalizeEvaluationResult<TQuestions extends EvaluateQuestions>(
  result: Experimental_EvaluationResult<TQuestions>,
): EvaluateResult<TQuestions> {
  return {
    answers: result.answers,
    providerMetadata: evaluationResultSchema.shape.providerMetadata.parse(
      result.providerMetadata ?? {},
    ),
    response: {
      headers: result.response.headers,
      id: result.response.id,
      modelId: result.response.modelId,
      timestamp: result.response.timestamp.toISOString(),
    },
    rounding: result.rounding,
    usage: result.usage,
    warnings: result.warnings,
  }
}