import { z } from "zod";

export const taskSchema = z.enum([
  "text-to-image",
  "image-to-image",
  "image-editing",
  "text-to-video",
  "image-to-video",
  "video-to-video",
  "video-editing",
  "speech-to-video",
  "lip-sync",
  "text-to-music",
  "text-to-speech",
  "speech-to-text",
  "audio-to-audio"
]);

export type KieTask = z.infer<typeof taskSchema>;

export const modelSchema = z.object({
  id: z.string(),
  name: z.string(),
  provider: z.string(),
  tasks: z.array(taskSchema),
  slug: z.string().optional(),
  price: z
    .object({
      unit: z.string().default("credits"),
      amount: z.number().optional()
    })
    .optional(),
  metadata: z.record(z.unknown()).optional(),
  description: z.string().optional(),
  status: z.enum(["active", "beta", "deprecated"]).default("active")
});

export type KieModel = z.infer<typeof modelSchema>;

export const jobStatusSchema = z.enum(["queued", "running", "succeeded", "failed", "cancelled"]);

export type KieJobStatus = z.infer<typeof jobStatusSchema>;

export const jobSchema = z.object({
  id: z.string(),
  modelId: z.string(),
  status: jobStatusSchema,
  createdAt: z.string(),
  updatedAt: z.string().optional(),
  resultUrl: z.string().url().optional(),
  error: z
    .object({
      message: z.string(),
      code: z.string().optional()
    })
    .optional(),
  metadata: z.record(z.unknown()).optional()
});

export type KieJob = z.infer<typeof jobSchema>;

export const queuePayloadSchema = z.object({
  modelId: z.string(),
  prompt: z.string().optional(),
  settings: z.record(z.unknown()).optional(),
  webhookUrl: z.string().url().optional(),
  assets: z
    .array(
      z.object({
        type: z.enum(["image", "video", "audio", "text"]),
        url: z.string().url(),
        role: z.string().optional()
      })
    )
    .optional()
});

export type QueuePayload = z.infer<typeof queuePayloadSchema>;
