import { citationSchema, type Citation } from '../types/citation';
import type { JSONObject } from '@ai-sdk/provider';
import {
  providerMetadataSchema,
  type ProviderMetadata,
} from '../types/provider-metadata';
import { jsonValueSchema } from '../types/json-value';
import type { FinishReason } from '../types/language-model';
import type {
  FileUIPart,
  InferUIMessageData,
  InferUIMessageMetadata,
  TextUIPart,
  UIDataTypes,
  UIMessage,
} from '../ui/ui-messages';
import type { ValueOf } from '../util/value-of';
import { z, type ZodType } from '../util/zod';
import { lazySchema, zodSchema } from '@ai-sdk/provider-utils';

const toolMetadataSchema: ZodType<JSONObject> = z.record(
  z.string(),
  z.optional(jsonValueSchema),
);

export const uiMessageChunkSchema = lazySchema(() =>
  zodSchema<UIMessageChunk>(
    z.union([
      z.looseObject({
        type: z.literal('text-start'),
        id: z.string(),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('text-delta'),
        id: z.string(),
        delta: z.string(),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('text-end'),
        id: z.string(),
        citations: z.optional(z.array(citationSchema)),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('error'),
        errorText: z.string(),
      }),
      z.looseObject({
        type: z.literal('tool-input-start'),
        toolCallId: z.string(),
        toolName: z.string(),
        inputEncoding: z.optional(z.enum(['json', 'text'])),
        providerExecuted: z.optional(z.boolean()),
        providerMetadata: z.optional(providerMetadataSchema),
        toolMetadata: z.optional(toolMetadataSchema),
        dynamic: z.optional(z.boolean()),
        title: z.optional(z.string()),
      }),
      z.looseObject({
        type: z.literal('tool-input-delta'),
        toolCallId: z.string(),
        inputTextDelta: z.string(),
      }),
      z.looseObject({
        type: z.literal('tool-input-available'),
        toolCallId: z.string(),
        toolName: z.string(),
        input: z.unknown(),
        providerExecuted: z.optional(z.boolean()),
        providerMetadata: z.optional(providerMetadataSchema),
        toolMetadata: z.optional(toolMetadataSchema),
        dynamic: z.optional(z.boolean()),
        title: z.optional(z.string()),
      }),
      z.looseObject({
        type: z.literal('tool-input-error'),
        toolCallId: z.string(),
        toolName: z.string(),
        input: z.unknown(),
        providerExecuted: z.optional(z.boolean()),
        providerMetadata: z.optional(providerMetadataSchema),
        toolMetadata: z.optional(toolMetadataSchema),
        dynamic: z.optional(z.boolean()),
        errorText: z.string(),
        title: z.optional(z.string()),
      }),
      z.looseObject({
        type: z.literal('tool-approval-request'),
        approvalId: z.string(),
        toolCallId: z.string(),
        approvalDescriptor: z.optional(z.unknown()),
        inputSchemaInput: z.optional(z.unknown()),
        reason: z.optional(z.string()),
        isAutomatic: z.optional(z.boolean()),
        signature: z.optional(z.string()),
      }),
      z.looseObject({
        type: z.literal('tool-approval-response'),
        approvalId: z.string(),
        approved: z.boolean(),
        reason: z.optional(z.string()),
        providerExecuted: z.optional(z.boolean()),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('tool-output-available'),
        toolCallId: z.string(),
        output: z.unknown(),
        providerExecuted: z.optional(z.boolean()),
        providerMetadata: z.optional(providerMetadataSchema),
        toolMetadata: z.optional(toolMetadataSchema),
        dynamic: z.optional(z.boolean()),
        preliminary: z.optional(z.boolean()),
      }),
      z.looseObject({
        type: z.literal('tool-output-error'),
        toolCallId: z.string(),
        errorText: z.string(),
        providerExecuted: z.optional(z.boolean()),
        providerMetadata: z.optional(providerMetadataSchema),
        toolMetadata: z.optional(toolMetadataSchema),
        dynamic: z.optional(z.boolean()),
      }),
      z.looseObject({
        type: z.literal('tool-output-denied'),
        toolCallId: z.string(),
      }),
      z.looseObject({
        type: z.literal('reasoning-start'),
        id: z.string(),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('reasoning-delta'),
        id: z.string(),
        delta: z.string(),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('reasoning-end'),
        id: z.string(),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('custom'),
        kind: z.pipe(
          z.string(),
          z.transform(value => value as `${string}.${string}`),
        ),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('source-url'),
        sourceId: z.string(),
        url: z.string(),
        title: z.optional(z.string()),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('source-document'),
        sourceId: z.string(),
        mediaType: z.string(),
        title: z.string(),
        filename: z.optional(z.string()),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('file'),
        url: z.string(),
        mediaType: z.string(),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('reasoning-file'),
        url: z.string(),
        mediaType: z.string(),
        providerMetadata: z.optional(providerMetadataSchema),
      }),
      z.looseObject({
        type: z.literal('steering-message'),
        id: z.string(),
        parts: z.array(
          z.union([
            z.looseObject({
              type: z.literal('text'),
              text: z.string(),
              providerMetadata: z.optional(providerMetadataSchema),
            }),
            z.looseObject({
              type: z.literal('file'),
              mediaType: z.string(),
              url: z.string(),
              filename: z.optional(z.string()),
              providerReference: z.optional(z.record(z.string(), z.string())),
              providerMetadata: z.optional(providerMetadataSchema),
            }),
          ]),
        ),
      }),
      z.looseObject({
        type: z.custom<`data-${string}`>(
          (value): value is `data-${string}` =>
            typeof value === 'string' && value.startsWith('data-'),
          { message: 'Type must start with "data-"' },
        ),
        id: z.optional(z.string()),
        data: z.unknown(),
        transient: z.optional(z.boolean()),
      }),
      z.looseObject({
        type: z.literal('start-step'),
      }),
      z.looseObject({
        type: z.literal('finish-step'),
      }),
      z.looseObject({
        type: z.literal('reset-step'),
      }),
      z.looseObject({
        type: z.literal('start'),
        messageId: z.optional(z.string()),
        messageMetadata: z.optional(z.unknown()),
      }),
      z.looseObject({
        type: z.literal('finish'),
        finishReason: z.optional(
          z.enum([
            'stop',
            'length',
            'content-filter',
            'tool-calls',
            'error',
            'interrupted',
            'other',
          ] as const satisfies readonly FinishReason[]),
        ),
        messageMetadata: z.optional(z.unknown()),
      }),
      z.looseObject({
        type: z.literal('abort'),
        reason: z.optional(z.string()),
      }),
      z.looseObject({
        type: z.literal('message-metadata'),
        messageMetadata: z.unknown(),
      }),
    ]),
  ),
);

export type DataUIMessageChunk<DATA_TYPES extends UIDataTypes> = ValueOf<{
  [NAME in keyof DATA_TYPES & string]: {
    type: `data-${NAME}`;
    id?: string;
    data: DATA_TYPES[NAME];
    transient?: boolean;
  };
}>;

export type UIMessageChunk<
  METADATA = unknown,
  DATA_TYPES extends UIDataTypes = UIDataTypes,
> =
  | {
      type: 'text-start';
      id: string;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'text-delta';
      delta: string;
      id: string;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'text-end';
      citations?: Array<Citation>;
      id: string;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'reasoning-start';
      id: string;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'reasoning-delta';
      id: string;
      delta: string;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'reasoning-end';
      id: string;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'custom';
      kind: `${string}.${string}`;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'error';
      errorText: string;
    }
  | {
      type: 'tool-input-available';
      toolCallId: string;
      toolName: string;
      input: unknown;
      providerExecuted?: boolean;
      providerMetadata?: ProviderMetadata;
      toolMetadata?: JSONObject;
      dynamic?: boolean;
      title?: string;
    }
  | {
      type: 'tool-input-error';
      toolCallId: string;
      toolName: string;
      input: unknown;
      providerExecuted?: boolean;
      providerMetadata?: ProviderMetadata;
      toolMetadata?: JSONObject;
      dynamic?: boolean;
      errorText: string;
      title?: string;
    }
  | {
      type: 'tool-approval-request';
      approvalId: string;
      toolCallId: string;
      approvalDescriptor?: unknown;
      inputSchemaInput?: unknown;
      reason?: string;
      isAutomatic?: boolean;
      signature?: string;
    }
  | {
      type: 'tool-approval-response';
      approvalId: string;
      approved: boolean;
      reason?: string;
      providerExecuted?: boolean;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'tool-output-available';
      toolCallId: string;
      output: unknown;
      providerExecuted?: boolean;
      providerMetadata?: ProviderMetadata;
      toolMetadata?: JSONObject;
      dynamic?: boolean;
      preliminary?: boolean;
    }
  | {
      type: 'tool-output-error';
      toolCallId: string;
      errorText: string;
      providerExecuted?: boolean;
      providerMetadata?: ProviderMetadata;
      toolMetadata?: JSONObject;
      dynamic?: boolean;
    }
  | {
      type: 'tool-output-denied';
      toolCallId: string;
    }
  | {
      type: 'tool-input-start';
      toolCallId: string;
      toolName: string;
      /** Encoding of the input deltas. Defaults to JSON when omitted. */
      inputEncoding?: 'json' | 'text';
      providerExecuted?: boolean;
      providerMetadata?: ProviderMetadata;
      toolMetadata?: JSONObject;
      dynamic?: boolean;
      title?: string;
    }
  | {
      type: 'tool-input-delta';
      toolCallId: string;
      inputTextDelta: string;
    }
  | {
      type: 'source-url';
      sourceId: string;
      url: string;
      title?: string;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'source-document';
      sourceId: string;
      mediaType: string;
      title: string;
      filename?: string;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'file';
      url: string;
      mediaType: string;
      providerMetadata?: ProviderMetadata;
    }
  | {
      type: 'reasoning-file';
      url: string;
      mediaType: string;
      providerMetadata?: ProviderMetadata;
    }
  | DataUIMessageChunk<DATA_TYPES>
  | {
      /**
       * Experimental. A user message that was added to the run with steering.
       */
      type: 'steering-message';
      id: string;
      parts: Array<TextUIPart | FileUIPart>;
    }
  | {
      type: 'start-step';
    }
  | {
      type: 'finish-step';
    }
  | {
      /**
       * Removes all message parts added during the current step.
       */
      type: 'reset-step';
    }
  | {
      type: 'start';
      messageId?: string;
      messageMetadata?: METADATA;
    }
  | {
      type: 'finish';
      finishReason?: FinishReason;
      messageMetadata?: METADATA;
    }
  | {
      type: 'abort';
      reason?: string;
    }
  | {
      type: 'message-metadata';
      messageMetadata: METADATA;
    };

export function isDataUIMessageChunk(
  chunk: UIMessageChunk,
): chunk is DataUIMessageChunk<UIDataTypes> {
  return chunk.type.startsWith('data-');
}

export type InferUIMessageChunk<T extends UIMessage> = UIMessageChunk<
  InferUIMessageMetadata<T>,
  InferUIMessageData<T>
>;