import * as z from "zod";
import { Id } from "./ids.js";
import { MediaFitSchema, MediaPresetIdSchema } from "./media-presets.js";
import { MediaTitleSchema } from "./media-video.js";

export const MEDIA_IMAGE_MAX_COUNT = 4;
export const MEDIA_VIDEO_MAX_SECONDS = 10;

const Prompt = z.string().trim().min(1).max(4000);
const NegativePrompt = z.string().trim().max(4000).optional().describe("What to avoid.");
const Seed = z
  .number()
  .int()
  .min(0)
  .max(Number.MAX_SAFE_INTEGER)
  .optional()
  .describe("Reuse a seed to reproduce a result.");
const Strength = z.number().min(0).max(1).optional();
const ASPECTS = "vertical 9:16, landscape 16:9, square 1:1, portrait 4:5, classic 4:3, cinema 21:9";
const aspect = (fallback: string) => MediaPresetIdSchema.describe(`${ASPECTS}. ${fallback}`);
const Side = z.number().int().min(64).max(2048);
const Pixels = z.number().int().min(0).max(16_384);

export const GenerateImageInputSchema = z
  .object({
    prompt: Prompt.describe("What the image shows."),
    negative_prompt: NegativePrompt,
    aspect: aspect("Default square, or the source image's aspect ratio.").optional(),
    width: Side.optional().describe("Exact width in pixels, with height; overrides aspect."),
    height: Side.optional(),
    count: z.number().int().min(1).max(MEDIA_IMAGE_MAX_COUNT).default(1),
    seed: Seed,
    source_artifact_id: Id.optional().describe("Image artifact to start from (image-to-image)."),
    strength: Strength.describe("With a source: how far to move away from it (default 0.6)."),
  })
  .refine((input) => (input.width === undefined) === (input.height === undefined), {
    message: "Set width and height together",
    path: ["width"],
  });
export type GenerateImageInput = z.infer<typeof GenerateImageInputSchema>;

export const GenerateVideoInputSchema = z.object({
  prompt: Prompt.describe("What happens in the clip: subject, motion, camera."),
  negative_prompt: NegativePrompt,
  aspect: aspect("Default landscape, or the start frame's aspect ratio.").optional(),
  duration_seconds: z.number().min(1).max(MEDIA_VIDEO_MAX_SECONDS).default(5),
  seed: Seed,
  source_artifact_id: Id.optional().describe("Image artifact to use as the first frame."),
});
export type GenerateVideoInput = z.infer<typeof GenerateVideoInputSchema>;

export const ImageFormatSchema = z.enum(["png", "jpeg", "webp"]);
export type ImageFormat = z.infer<typeof ImageFormatSchema>;

/** Edits apply in this order: crop, rotate, flip, resize, extend, repaint, format. */
export const EditImageInputSchema = z
  .object({
    artifact_id: Id.describe("Image artifact to edit."),
    title: MediaTitleSchema.optional().describe(
      "A few words in the user's language for the card and file name.",
    ),
    crop: z
      .object({ left: Pixels, top: Pixels, width: Pixels.min(1), height: Pixels.min(1) })
      .optional()
      .describe("Area to keep, in source pixels."),
    rotate: z.literal([90, 180, 270]).optional().describe("Clockwise, in degrees."),
    flip: z.enum(["horizontal", "vertical"]).optional(),
    resize: z
      .object({
        aspect: aspect("Resizes to the format's exact size.").optional(),
        width: z.number().int().min(1).max(8192).optional(),
        height: z.number().int().min(1).max(8192).optional(),
        fit: MediaFitSchema.default("crop").describe(
          "How a different aspect ratio fills the size: crop, pad, or blur (blurred background).",
        ),
      })
      .refine((resize) => resize.aspect || resize.width || resize.height, {
        message: "Set aspect, width or height",
      })
      .optional()
      .describe("Width or height alone keeps the aspect ratio."),
    extend: z
      .object({
        top: Pixels.default(0),
        right: Pixels.default(0),
        bottom: Pixels.default(0),
        left: Pixels.default(0),
        color: z
          .string()
          .regex(/^#[0-9a-fA-F]{6}$/)
          .optional()
          .describe("#RRGGBB; transparent when omitted and the format allows it."),
      })
      .optional()
      .describe("Borders to add, in pixels."),
    prompt: Prompt.optional().describe("Repaint the image with this prompt."),
    strength: Strength.describe("How far the repaint moves away from the image (default 0.6)."),
    format: ImageFormatSchema.optional().describe("Default keeps the source format."),
  })
  .refine(
    (input) =>
      input.crop ||
      input.rotate ||
      input.flip ||
      input.resize ||
      input.extend ||
      input.prompt ||
      input.format,
    { message: "Nothing to edit" },
  );
export type EditImageInput = z.infer<typeof EditImageInputSchema>;