import { z } from 'zod';


// text, e.g.: { 'type': 'text', 'text': 'Hello, Claude' }
const anthropicWireTextBlockSchema = z.object({
  type: z.literal('text'),
  text: z.string(),
});

// image, e.g.: { 'type': 'image', 'source': { 'type': 'base64', 'media_type': 'image/jpeg', 'data': '/9j/4AAQSkZJRg...' } }
const anthropicWireImageBlockSchema = z.object({
  type: z.literal('image'),
  source: z.object({
    type: z.enum(['base64']),
    media_type: z.enum(['image/jpeg', 'image/png', 'image/gif', 'image/webp']),
    data: z.string(),
  }),
});

const anthropicWireMessagesSchema = z.array(
  z.object({
    role: z.enum(['user', 'assistant']),
    // NOTE: could be a string or an array of text/image blocks, but for a better implementation
    //       we will assume it's always an array
    // content: z.union([
    //   z.array(z.union([anthropicWireTextBlockSchema, anthropicWireImageBlockSchema])),
    //   z.string(),
    // ]),
    content: z.array(
      z.union([
        anthropicWireTextBlockSchema,
        anthropicWireImageBlockSchema,
      ]),
    ),
  }),
);

export const anthropicWireMessagesRequestSchema = z.object({
  model: z.string(),

  /**
   * If you want to include a system prompt, you can use the top-level system parameter — there is no "system" role for input messages in the Messages API.
   */
  system: z.string().optional(),

  /**
   * (required) Input messages. - operates on alternating user and assistant conversational turns - the first message must always use the user role
   * If the final message uses the assistant role, the response content will continue immediately from the content in that message.
   * This can be used to constrain part of the model's response.
   */
  messages: anthropicWireMessagesSchema.refine(
    (messages) => {

      // Ensure the first message uses the user role
      if (messages.length === 0 || messages[0].role !== 'user')
        return false;

      // Ensure messages alternate between user and assistant roles
      for (let i = 1; i < messages.length; i++)
        if (messages[i].role === messages[i - 1].role)
          return false;

      return true;
    },
    { message: `messages must alternate between User and Assistant roles, starting with the User role` },
  ),

  /**
   * (required) The maximum number of tokens to generate before stopping.
   */
  max_tokens: z.number(),


  /**
   * (optional) Metadata to include with the request.
   * user_id: This should be a uuid, hash value, or other opaque identifier.
   */
  metadata: z.object({
    user_id: z.string().optional(),
  }).optional(),

  /**
   * Custom text sequences that will cause the model to stop generating.
   */
  stop_sequences: z.array(z.string()).optional(),

  /**
   * Whether to incrementally stream the response using server-sent events. Default: false
   */
  stream: z.boolean().optional(),

  /**
   * Defaults to 1.0. Ranges from 0.0 to 1.0. Use temperature closer to 0.0 for analytical / multiple choice, and closer to 1.0 for creative and generative tasks.
   */
  temperature: z.number().optional(),

  /**
   * Use nucleus sampling.
   * Recommended for advanced use cases only. You usually only need to use temperature.
   */
  top_p: z.number().optional(),

  /**
   * Only sample from the top K options for each subsequent token.
   * Recommended for advanced use cases only. You usually only need to use temperature.
   */
  top_k: z.number().optional(),
});
export type AnthropicWireMessagesRequest = z.infer<typeof anthropicWireMessagesRequestSchema>;


export const anthropicWireMessagesResponseSchema = z.object({
  // Unique object identifier.
  id: z.string(),

  // For Messages, this is always "message".
  type: z.literal('message'),
  // Conversational role of the generated message. This will always be "assistant".
  role: z.literal('assistant'),
  /**
   * Content generated by the model.
   * This is an array of content blocks, each of which has a type that determines its shape. Currently, the only type in responses is "text".
   */
  content: z.array(anthropicWireTextBlockSchema),

  // The model that handled the request.
  model: z.string(),

  /**
   * This may be one the following values:
   *
   * "end_turn": the model reached a natural stopping point
   * "max_tokens": we exceeded the requested max_tokens or the model's maximum
   * "stop_sequence": one of your provided custom stop_sequences was generated
   * Note that these values are different than those in /v1/complete, where end_turn and stop_sequence were not differentiated.
   *
   * In non-streaming mode this value is always non-null. In streaming mode, it is null in the message_start event and non-null otherwise.
   */
  stop_reason: z.enum(['end_turn', 'max_tokens', 'stop_sequence']).nullable(),

  // Which custom stop sequence was generated, if any.
  stop_sequence: z.string().nullable(),

  // Billing and rate-limit usage.
  usage: z.object({
    input_tokens: z.number(),
    output_tokens: z.number(),
  }),

});
export type AnthropicWireMessagesResponse = z.infer<typeof anthropicWireMessagesResponseSchema>;
