import z from 'zod';
import { query, queryOrigin } from './QueryService';
import { snippet } from './SnippetService';

export const zPortalActions = z.enum([
  'Open',
  'Close',
  'Focus',
]);

export type PortalActions = z.infer<typeof zPortalActions>;
/** ******************************************************************************
*  Semantic Search
******************************************************************************* */

export const semanticSearchParams = z.object({
  botId: z.string(),
  sourceIds: z.array(z.string()),
  syncIds: z.array(z.string()).optional(),
  query: z.string(),
  topK: z.number().optional(),
  distinctId: z.string(),
  origin: queryOrigin.nullable(),
});

export const semanticSearchResult = z.object({
  snippets: z.array(snippet),
});

export type SemanticSearchParams = z.infer<typeof semanticSearchParams>;
export type SemanticSearchResult = z.infer<typeof semanticSearchResult>;

/** ******************************************************************************
*  Create Generation
******************************************************************************* */

export const createGenerationParams = z.object({
  botId: z.string(),
  sourceIds: z.array(z.string()),
  query: z.string(),
  distinctId: z.string(),
  origin: queryOrigin.optional(),
});

export const createGenerationResponse = z.object({
  query: query.optional(),
  snippets: z.array(snippet).optional(),
  error: z.string().optional(),
});

export type CreateGenerationParams = z.infer<typeof createGenerationParams>;
export type CreateGenerationResponse = z.infer<typeof createGenerationResponse>;