import {
  GenkitError,
  OperationSchema,
  action,
  backgroundAction,
  defineAction,
  registerBackgroundAction,
  z
} from "@genkit-ai/core";
import { logger } from "@genkit-ai/core/logging";
import { toJsonSchema } from "@genkit-ai/core/schema";
import { performance } from "node:perf_hooks";
import {
  CustomPartSchema,
  DataPartSchema,
  MediaPartSchema,
  TextPartSchema,
  ToolRequestPartSchema,
  ToolResponsePartSchema
} from "./document.js";
import {
  GenerateRequestSchema,
  GenerateResponseSchema
} from "./model-types.js";
import {
  augmentWithContext,
  simulateConstrainedGeneration
} from "./model/middleware.js";
import { defineGenerateAction } from "./generate/action.js";
export * from "./model-types.js";
function model(options, runner) {
  const act = action(modelActionOptions(options), (input, ctx) => {
    const startTimeMs = performance.now();
    return runner(input, ctx).then((response) => {
      const timedResponse = {
        ...response,
        latencyMs: performance.now() - startTimeMs
      };
      return timedResponse;
    });
  });
  Object.assign(act, {
    __configSchema: options.configSchema || z.unknown()
  });
  return act;
}
function modelActionOptions(options) {
  const label = options.label || options.name;
  const middleware = getModelMiddleware(options);
  return {
    actionType: "model",
    name: options.name,
    description: label,
    inputSchema: GenerateRequestSchema,
    outputSchema: GenerateResponseSchema,
    metadata: {
      model: {
        label,
        customOptions: options.configSchema ? toJsonSchema({ schema: options.configSchema }) : void 0,
        versions: options.versions,
        supports: options.supports
      }
    },
    use: middleware
  };
}
function defineModel(registry, options, runner) {
  const act = defineAction(
    registry,
    modelActionOptions(options),
    (input, ctx) => {
      const startTimeMs = performance.now();
      const secondParam = options.apiVersion === "v2" ? ctx : ctx.streamingRequested ? ctx.sendChunk : void 0;
      return runner(input, secondParam).then((response) => {
        const timedResponse = {
          ...response,
          latencyMs: performance.now() - startTimeMs
        };
        return timedResponse;
      });
    }
  );
  Object.assign(act, {
    __configSchema: options.configSchema || z.unknown()
  });
  return act;
}
function defineBackgroundModel(registry, options) {
  const act = backgroundModel(options);
  registerBackgroundAction(registry, act);
  return act;
}
function backgroundModel(options) {
  const label = options.label || options.name;
  const middleware = getModelMiddleware(options);
  const act = backgroundAction({
    actionType: "background-model",
    name: options.name,
    description: label,
    inputSchema: GenerateRequestSchema,
    outputSchema: GenerateResponseSchema,
    metadata: {
      model: {
        label,
        customOptions: options.configSchema ? toJsonSchema({ schema: options.configSchema }) : void 0,
        versions: options.versions,
        supports: options.supports
      }
    },
    use: middleware,
    async start(request) {
      const startTimeMs = performance.now();
      const response = await options.start(request);
      Object.assign(response, {
        latencyMs: performance.now() - startTimeMs
      });
      return response;
    },
    async check(op) {
      return options.check(op);
    },
    cancel: options.cancel ? async (op) => {
      if (!options.cancel) {
        throw new GenkitError({
          status: "UNIMPLEMENTED",
          message: "cancel not implemented"
        });
      }
      return options.cancel(op);
    } : void 0
  });
  Object.assign(act, {
    __configSchema: options.configSchema || z.unknown()
  });
  return act;
}
function getModelMiddleware(options) {
  const middleware = options.use || [];
  if (!options?.supports?.context) middleware.push(augmentWithContext());
  const constratedSimulator = simulateConstrainedGeneration();
  middleware.push((req, next) => {
    if (!options?.supports?.constrained || options?.supports?.constrained === "none" || options?.supports?.constrained === "no-tools" && (req.tools?.length ?? 0) > 0) {
      return constratedSimulator(req, next);
    }
    return next(req);
  });
  return middleware;
}
function modelActionMetadata({
  name,
  info,
  configSchema,
  background
}) {
  return {
    actionType: background ? "background-model" : "model",
    name,
    inputJsonSchema: toJsonSchema({ schema: GenerateRequestSchema }),
    outputJsonSchema: background ? toJsonSchema({ schema: OperationSchema }) : toJsonSchema({ schema: GenerateResponseSchema }),
    metadata: {
      model: {
        ...info,
        customOptions: configSchema ? toJsonSchema({ schema: configSchema }) : void 0
      }
    }
  };
}
function modelRef(options) {
  let name = options.name;
  if (options.namespace && !name.startsWith(options.namespace + "/")) {
    name = `${options.namespace}/${name}`;
  }
  const ref = {
    ...options,
    name
  };
  ref.withConfig = (cfg) => {
    return modelRef({
      ...options,
      name,
      config: cfg
    });
  };
  ref.withVersion = (version) => {
    return modelRef({
      ...options,
      name,
      version
    });
  };
  return ref;
}
function getBasicUsageStats(input, response) {
  const inputCounts = getPartCounts(input.flatMap((md) => md.content));
  const outputCounts = getPartCounts(
    Array.isArray(response) ? response.flatMap((c) => c.message.content) : response.content
  );
  return {
    inputCharacters: inputCounts.characters,
    inputImages: inputCounts.images,
    inputVideos: inputCounts.videos,
    inputAudioFiles: inputCounts.audio,
    outputCharacters: outputCounts.characters,
    outputImages: outputCounts.images,
    outputVideos: outputCounts.videos,
    outputAudioFiles: outputCounts.audio
  };
}
function getPartCounts(parts) {
  return parts.reduce(
    (counts, part) => {
      const isImage = part.media?.contentType?.startsWith("image") || part.media?.url?.startsWith("data:image");
      const isVideo = part.media?.contentType?.startsWith("video") || part.media?.url?.startsWith("data:video");
      const isAudio = part.media?.contentType?.startsWith("audio") || part.media?.url?.startsWith("data:audio");
      return {
        characters: counts.characters + (part.text?.length || 0),
        images: counts.images + (isImage ? 1 : 0),
        videos: counts.videos + (isVideo ? 1 : 0),
        audio: counts.audio + (isAudio ? 1 : 0)
      };
    },
    { characters: 0, images: 0, videos: 0, audio: 0 }
  );
}
async function resolveModel(registry, model2, options) {
  let out;
  let modelId;
  if (!model2) {
    model2 = await registry.lookupValue("defaultModel", "defaultModel");
  }
  if (!model2) {
    throw new GenkitError({
      status: "INVALID_ARGUMENT",
      message: "Must supply a `model` to `generate()` calls."
    });
  }
  if (typeof model2 === "string") {
    modelId = model2;
    out = { modelAction: await lookupModel(registry, model2) };
  } else if (model2.hasOwnProperty("__action")) {
    modelId = model2.__action.name;
    out = { modelAction: model2 };
  } else {
    const ref = model2;
    modelId = ref.name;
    out = {
      modelAction: await lookupModel(registry, ref.name),
      config: {
        ...ref.config
      },
      version: ref.version
    };
  }
  if (!out.modelAction) {
    throw new GenkitError({
      status: "NOT_FOUND",
      message: `Model '${modelId}' not found`
    });
  }
  if (options?.warnDeprecated && out.modelAction.__action.metadata?.model?.stage === "deprecated") {
    logger.warn(
      `Model '${out.modelAction.__action.name}' is deprecated and may be removed in a future release.`
    );
  }
  return out;
}
async function lookupModel(registry, model2) {
  return await registry.lookupAction(`/model/${model2}`) || await registry.lookupAction(`/background-model/${model2}`);
}
export {
  CustomPartSchema,
  DataPartSchema,
  MediaPartSchema,
  TextPartSchema,
  ToolRequestPartSchema,
  ToolResponsePartSchema,
  backgroundModel,
  defineBackgroundModel,
  defineGenerateAction,
  defineModel,
  getBasicUsageStats,
  model,
  modelActionMetadata,
  modelRef,
  resolveModel,
  simulateConstrainedGeneration
};
//# sourceMappingURL=model.mjs.map