import { z } from 'zod';
import { SpecifyDesignTokenFormat } from '@specifyapp/specify-design-token-format';

export const textOutputSchema = z
  .object({
    type: z.literal('text'),
    text: z.string(),
  })
  .strict();
export type TextOutput = {
  type: 'text';
  text: string;
};

export const sdtfOutputSchema = z
  .object({
    type: z.literal('SDTF'),
    graph: z.custom<SpecifyDesignTokenFormat>(value => value),
  })
  .strict();
export type SDTFOutput = {
  type: 'SDTF';
  graph: SpecifyDesignTokenFormat;
};

export const jsonOutputSchema = z
  .object({
    type: z.literal('JSON'),
    json: z.union([
      z.record(z.string(), z.unknown()),
      z.array(z.unknown()),
      z.string(),
      z.number(),
      z.boolean(),
      z.null(),
    ]),
  })
  .strict();
export type JSONOutput = {
  type: 'JSON';
  json: unknown;
};

export const filesOutputFileSchema = z
  .object({
    path: z.string(),
    content: z.union([
      z.object({
        type: z.literal('text'),
        text: z.string(),
      }),
      z.object({
        type: z.literal('url'),
        url: z.string(),
      }),
    ]),
  })
  .strict();
export const filesOutputSchema = z
  .object({
    type: z.literal('files'),
    files: z.array(filesOutputFileSchema),
  })
  .strict();
export type FilesOutput = {
  type: 'files';
  files: Array<{
    path: string; // path as in configuration
    content: { type: 'text'; text: string } | { type: 'url'; url: string };
  }>;
};
export type ParserOutput = TextOutput | SDTFOutput | JSONOutput | FilesOutput;
export type ParserOutputType = ParserOutput['type'];

export const parserOutputSchema = z.discriminatedUnion('type', [
  textOutputSchema,
  sdtfOutputSchema,
  jsonOutputSchema,
  filesOutputSchema,
]);