import { z } from 'zod';

// ── Shared sub-schemas ──────────────────────────────────────────

export const ExtractionExampleSchema = z.object({
  text: z.string().min(1),
  extractions: z.array(
    z.object({
      extraction_class: z.string().min(1),
      extraction_text: z.string().min(1),
      attributes: z.record(z.string()).optional(),
    })
  ),
});

export const ExtractionResultSchema = z.object({
  extraction_class: z.string(),
  extraction_text: z.string(),
  attributes: z.record(z.string()).optional(),
  start_offset: z.number().optional(),
  end_offset: z.number().optional(),
});

// ── Request schemas ─────────────────────────────────────────────

export const ExtractRequestSchema = z.object({
  text: z.string().min(1, 'text is required').max(50_000, 'text exceeds 50,000 character limit'),
  taskId: z.string().optional(),
  taskPrompt: z.string().optional(),
  examples: z.array(ExtractionExampleSchema).optional(),
  modelId: z.string().optional(),
  options: z
    .object({
      extractionPasses: z.number().int().min(1).max(5).optional(),
      maxWorkers: z.number().int().min(1).max(50).optional(),
      maxCharBuffer: z.number().int().min(100).max(10_000).optional(),
    })
    .optional(),
  productId: z.string().optional(),
});

export const BatchExtractRequestSchema = z.object({
  inputs: z
    .array(
      z.object({
        text: z.string().min(1).max(50_000),
        taskId: z.string().optional(),
        taskPrompt: z.string().optional(),
      })
    )
    .min(1)
    .max(50),
  examples: z.array(ExtractionExampleSchema).optional(),
  modelId: z.string().optional(),
  productId: z.string().optional(),
});

// ── Response schemas ────────────────────────────────────────────

export const ExtractResponseSchema = z.object({
  extractions: z.array(ExtractionResultSchema),
  metadata: z.object({
    modelId: z.string(),
    durationMs: z.number(),
    tokenCount: z.number().optional(),
    charCount: z.number(),
  }),
  requestId: z.string().optional(),
});

// ── TypeScript types ────────────────────────────────────────────

export type ExtractionExample = z.infer<typeof ExtractionExampleSchema>;
export type ExtractionResult = z.infer<typeof ExtractionResultSchema>;
export type ExtractRequest = z.infer<typeof ExtractRequestSchema>;
export type BatchExtractRequest = z.infer<typeof BatchExtractRequestSchema>;
export type ExtractResponse = z.infer<typeof ExtractResponseSchema>;
