import {
  TypeValidationError,
  type JSONSchema7,
  type JSONValue,
  type LanguageModelV4CallOptions,
} from '@ai-sdk/provider';
import {
  asSchema,
  resolve,
  safeParseJSON,
  safeValidateTypes,
  type FlexibleSchema,
} from '@ai-sdk/provider-utils';
import { InvalidArgumentError } from '../error/invalid-argument-error';
import { NoObjectGeneratedError } from '../error/no-object-generated-error';
import type { FinishReason } from '../types/language-model';
import type { LanguageModelResponseMetadata } from '../types/language-model-response-metadata';
import type { LanguageModelUsage } from '../types/usage';
import type { DeepPartial } from '../util/deep-partial';
import { parsePartialJson } from '../util/parse-partial-json';
import type { EnrichedStreamPart } from './stream-text';

export interface Output<OUTPUT = any, PARTIAL = any, ELEMENT = any> {
  /**
   * The name of the output mode.
   */
  name: string;

  /**
   * The response format to use for the model.
   */
  responseFormat: PromiseLike<LanguageModelV4CallOptions['responseFormat']>;

  /**
   * Parses the complete output of the model.
   */
  parseCompleteOutput(
    options: { text: string },
    context: {
      response: Omit<LanguageModelResponseMetadata, 'messages' | 'body'>;
      usage: LanguageModelUsage;
      finishReason: FinishReason;
    },
  ): Promise<OUTPUT>;

  /**
   * Parses the partial output of the model.
   */
  parsePartialOutput(options: {
    text: string;
  }): Promise<{ partial: PARTIAL } | undefined>;

  /**
   * Creates a stream transform that emits individual elements as they complete.
   */
  createElementStreamTransform():
    | TransformStream<EnrichedStreamPart<any, PARTIAL>, ELEMENT>
    | undefined;
}

/**
 * Output specification for text generation.
 * This is the default output mode that generates plain text.
 *
 * @returns An output specification for generating text.
 */
export const text = (): Output<string, string, never> => ({
  name: 'text',
  responseFormat: Promise.resolve({ type: 'text' }),

  async parseCompleteOutput({ text }: { text: string }) {
    return text;
  },

  async parsePartialOutput({ text }: { text: string }) {
    return { partial: text };
  },

  createElementStreamTransform() {
    return undefined;
  },
});

/**
 * Output specification for typed object generation using schemas.
 * When the model generates a text response, it will return an object that matches the schema.
 *
 * @param schema - The schema of the object to generate.
 * @param name - Optional name of the output that should be generated. Used by some providers for additional LLM guidance, e.g. via tool or schema name.
 * @param description - Optional description of the output that should be generated. Used by some providers for additional LLM guidance, e.g. via tool or schema description.
 *
 * @returns An output specification for generating objects with the specified schema.
 */
export const object = <OBJECT>({
  schema: inputSchema,
  name,
  description,
}: {
  schema: FlexibleSchema<OBJECT>;
  /**
   * Optional name of the output that should be generated.
   * Used by some providers for additional LLM guidance, e.g. via tool or schema name.
   */
  name?: string;
  /**
   * Optional description of the output that should be generated.
   * Used by some providers for additional LLM guidance, e.g. via tool or schema description.
   */
  description?: string;
}): Output<OBJECT, DeepPartial<OBJECT>, never> => {
  const schema = asSchema(inputSchema);

  return {
    name: 'object',

    responseFormat: resolve(schema.jsonSchema).then(jsonSchema => ({
      type: 'json' as const,
      schema: jsonSchema,
      ...(name != null && { name }),
      ...(description != null && { description }),
    })),

    async parseCompleteOutput(
      { text }: { text: string },
      context: {
        response: LanguageModelResponseMetadata;
        usage: LanguageModelUsage;
        finishReason: FinishReason;
      },
    ) {
      const parseResult = await safeParseJSON({ text });

      if (!parseResult.success) {
        throw new NoObjectGeneratedError({
          message: 'No object generated: could not parse the response.',
          cause: parseResult.error,
          text,
          response: context.response,
          usage: context.usage,
          finishReason: context.finishReason,
        });
      }

      const validationResult = await safeValidateTypes({
        value: parseResult.value,
        schema,
      });

      if (!validationResult.success) {
        throw new NoObjectGeneratedError({
          message: 'No object generated: response did not match schema.',
          cause: validationResult.error,
          text,
          response: context.response,
          usage: context.usage,
          finishReason: context.finishReason,
        });
      }

      return validationResult.value;
    },

    async parsePartialOutput({ text }: { text: string }) {
      const result = await parsePartialJson(text);

      switch (result.state) {
        case 'failed-parse':
        case 'undefined-input': {
          return undefined;
        }

        case 'repaired-parse':
        case 'successful-parse': {
          return {
            // Note: currently no validation of partial results:
            partial: result.value as DeepPartial<OBJECT>,
          };
        }
      }
    },

    createElementStreamTransform() {
      return undefined;
    },
  };
};

/**
 * Output specification for array generation.
 * When the model generates a text response, it will return an array of elements.
 *
 * @param element - The schema of the array elements to generate.
 * @param minItems - Optional minimum number of elements to generate.
 * @param maxItems - Optional maximum number of elements to generate.
 * @param name - Optional name of the output that should be generated. Used by some providers for additional LLM guidance, e.g. via tool or schema name.
 * @param description - Optional description of the output that should be generated. Used by some providers for additional LLM guidance, e.g. via tool or schema description.
 *
 * @returns An output specification for generating an array of elements.
 */
export const array = <ELEMENT>({
  element: inputElementSchema,
  minItems,
  maxItems,
  name,
  description,
}: {
  element: FlexibleSchema<ELEMENT>;
  /**
   * Optional minimum number of elements to generate.
   */
  minItems?: number;
  /**
   * Optional maximum number of elements to generate.
   */
  maxItems?: number;
  /**
   * Optional name of the output that should be generated.
   * Used by some providers for additional LLM guidance, e.g. via tool or schema name.
   */
  name?: string;
  /**
   * Optional description of the output that should be generated.
   * Used by some providers for additional LLM guidance, e.g. via tool or schema description.
   */
  description?: string;
}): Output<Array<ELEMENT>, Array<ELEMENT>, ELEMENT> => {
  validateArrayBound({ name: 'minItems', value: minItems });
  validateArrayBound({ name: 'maxItems', value: maxItems });

  if (minItems != null && maxItems != null && minItems > maxItems) {
    throw new InvalidArgumentError({
      parameter: 'minItems',
      value: minItems,
      message: 'minItems must be less than or equal to maxItems',
    });
  }

  const elementSchema = asSchema(inputElementSchema);

  return {
    name: 'array',

    // JSON schema that describes an array of elements:
    responseFormat: resolve(elementSchema.jsonSchema).then(jsonSchema => {
      // keep root-level definitions available to root-relative references:
      const {
        $schema: _$schema,
        definitions,
        $defs,
        ...itemSchema
      } = jsonSchema as JSONSchema7 & {
        $defs?: JSONSchema7['definitions'];
      };

      return {
        type: 'json' as const,
        schema: {
          $schema: 'http://json-schema.org/draft-07/schema#',
          ...(definitions != null && { definitions }),
          ...($defs != null && { $defs }),
          type: 'object',
          properties: {
            elements: {
              type: 'array',
              items: itemSchema,
              ...(minItems != null && { minItems }),
              ...(maxItems != null && { maxItems }),
            },
          },
          required: ['elements'],
          additionalProperties: false,
        },
        ...(name != null && { name }),
        ...(description != null && { description }),
      };
    }),

    async parseCompleteOutput(
      { text }: { text: string },
      context: {
        response: LanguageModelResponseMetadata;
        usage: LanguageModelUsage;
        finishReason: FinishReason;
      },
    ) {
      const parseResult = await safeParseJSON({ text });

      if (!parseResult.success) {
        throw new NoObjectGeneratedError({
          message: 'No object generated: could not parse the response.',
          cause: parseResult.error,
          text,
          response: context.response,
          usage: context.usage,
          finishReason: context.finishReason,
        });
      }

      const outerValue = parseResult.value;

      if (
        outerValue == null ||
        typeof outerValue !== 'object' ||
        !('elements' in outerValue) ||
        !Array.isArray(outerValue.elements)
      ) {
        throw new NoObjectGeneratedError({
          message: 'No object generated: response did not match schema.',
          cause: new TypeValidationError({
            value: outerValue,
            cause: 'response must be an object with an elements array',
          }),
          text,
          response: context.response,
          usage: context.usage,
          finishReason: context.finishReason,
        });
      }

      const lengthValidationError = getArrayLengthValidationError({
        value: outerValue.elements,
        minItems,
        maxItems,
      });

      if (lengthValidationError != null) {
        throw new NoObjectGeneratedError({
          message: 'No object generated: response did not match schema.',
          cause: lengthValidationError,
          text,
          response: context.response,
          usage: context.usage,
          finishReason: context.finishReason,
        });
      }

      const validatedElements: Array<ELEMENT> = [];
      for (const element of outerValue.elements) {
        const validationResult = await safeValidateTypes({
          value: element,
          schema: elementSchema,
        });

        if (!validationResult.success) {
          throw new NoObjectGeneratedError({
            message: 'No object generated: response did not match schema.',
            cause: validationResult.error,
            text,
            response: context.response,
            usage: context.usage,
            finishReason: context.finishReason,
          });
        }

        validatedElements.push(validationResult.value);
      }

      return validatedElements;
    },

    async parsePartialOutput({ text }: { text: string }) {
      const result = await parsePartialJson(text);

      switch (result.state) {
        case 'failed-parse':
        case 'undefined-input': {
          return undefined;
        }

        case 'repaired-parse':
        case 'successful-parse': {
          const outerValue = result.value;

          // no parsable elements array
          if (
            outerValue == null ||
            typeof outerValue !== 'object' ||
            !('elements' in outerValue) ||
            !Array.isArray(outerValue.elements)
          ) {
            return undefined;
          }

          const rawElements =
            result.state === 'repaired-parse' && outerValue.elements.length > 0
              ? outerValue.elements.slice(0, -1)
              : outerValue.elements;

          const parsedElements: Array<ELEMENT> = [];
          for (const rawElement of rawElements) {
            const validationResult = await safeValidateTypes({
              value: rawElement,
              schema: elementSchema,
            });

            if (validationResult.success) {
              parsedElements.push(validationResult.value);
            }
          }

          return { partial: parsedElements };
        }
      }
    },

    createElementStreamTransform() {
      let publishedElements = 0;

      return new TransformStream<
        EnrichedStreamPart<any, Array<ELEMENT>>,
        ELEMENT
      >({
        transform({ partialOutput }, controller) {
          if (partialOutput != null) {
            // Only enqueue new elements that haven't been published yet
            for (
              ;
              publishedElements < partialOutput.length;
              publishedElements++
            ) {
              if (maxItems != null && publishedElements >= maxItems) {
                controller.error(
                  getArrayLengthValidationError({
                    value: partialOutput,
                    maxItems,
                  }),
                );
                return;
              }

              controller.enqueue(partialOutput[publishedElements]);
            }
          }
        },
      });
    },
  };
};

function validateArrayBound({
  name,
  value,
}: {
  name: 'minItems' | 'maxItems';
  value: number | undefined;
}) {
  if (value == null) {
    return;
  }

  if (!Number.isInteger(value)) {
    throw new InvalidArgumentError({
      parameter: name,
      value,
      message: `${name} must be an integer`,
    });
  }

  if (value < 0) {
    throw new InvalidArgumentError({
      parameter: name,
      value,
      message: `${name} must be greater than or equal to 0`,
    });
  }
}

function getArrayLengthValidationError({
  value,
  minItems,
  maxItems,
}: {
  value: Array<unknown>;
  minItems?: number;
  maxItems?: number;
}): TypeValidationError | undefined {
  if (minItems != null && value.length < minItems) {
    return new TypeValidationError({
      value,
      cause: `elements array must contain at least ${minItems} items`,
    });
  }

  if (maxItems != null && value.length > maxItems) {
    return new TypeValidationError({
      value,
      cause: `elements array must contain at most ${maxItems} items`,
    });
  }
}

/**
 * Output specification for choice generation.
 * When the model generates a text response, it will return a one of the choice options.
 *
 * @param options - The available choices.
 * @param name - Optional name of the output that should be generated. Used by some providers for additional LLM guidance, e.g. via tool or schema name.
 * @param description - Optional description of the output that should be generated. Used by some providers for additional LLM guidance, e.g. via tool or schema description.
 *
 * @returns An output specification for generating a choice.
 */
export const choice = <CHOICE extends string>({
  options: choiceOptions,
  name,
  description,
}: {
  options: Array<CHOICE>;
  /**
   * Optional name of the output that should be generated.
   * Used by some providers for additional LLM guidance, e.g. via tool or schema name.
   */
  name?: string;
  /**
   * Optional description of the output that should be generated.
   * Used by some providers for additional LLM guidance, e.g. via tool or schema description.
   */
  description?: string;
}): Output<CHOICE, CHOICE, never> => {
  return {
    name: 'choice',

    // JSON schema that describes an enumeration:
    responseFormat: Promise.resolve({
      type: 'json',
      schema: {
        $schema: 'http://json-schema.org/draft-07/schema#',
        type: 'object',
        properties: {
          result: { type: 'string', enum: choiceOptions },
        },
        required: ['result'],
        additionalProperties: false,
      },
      ...(name != null && { name }),
      ...(description != null && { description }),
    } as const),

    async parseCompleteOutput(
      { text }: { text: string },
      context: {
        response: LanguageModelResponseMetadata;
        usage: LanguageModelUsage;
        finishReason: FinishReason;
      },
    ) {
      const parseResult = await safeParseJSON({ text });

      if (!parseResult.success) {
        throw new NoObjectGeneratedError({
          message: 'No object generated: could not parse the response.',
          cause: parseResult.error,
          text,
          response: context.response,
          usage: context.usage,
          finishReason: context.finishReason,
        });
      }

      const outerValue = parseResult.value;

      if (
        outerValue == null ||
        typeof outerValue !== 'object' ||
        !('result' in outerValue) ||
        typeof outerValue.result !== 'string' ||
        !choiceOptions.includes(outerValue.result as any)
      ) {
        throw new NoObjectGeneratedError({
          message: 'No object generated: response did not match schema.',
          cause: new TypeValidationError({
            value: outerValue,
            cause: 'response must be an object that contains a choice value.',
          }),
          text,
          response: context.response,
          usage: context.usage,
          finishReason: context.finishReason,
        });
      }

      return outerValue.result as CHOICE;
    },

    async parsePartialOutput({ text }: { text: string }) {
      const result = await parsePartialJson(text);

      switch (result.state) {
        case 'failed-parse':
        case 'undefined-input': {
          return undefined;
        }

        case 'repaired-parse':
        case 'successful-parse': {
          const outerValue = result.value;

          if (
            outerValue == null ||
            typeof outerValue !== 'object' ||
            !('result' in outerValue) ||
            typeof outerValue.result !== 'string'
          ) {
            return undefined;
          }

          // list of potential matches.
          const potentialMatches = choiceOptions.filter(choiceOption =>
            choiceOption.startsWith(outerValue.result as string),
          );

          if (result.state === 'successful-parse') {
            // successful parse: exact choice value
            return potentialMatches.includes(outerValue.result as any)
              ? { partial: outerValue.result as CHOICE }
              : undefined;
          } else {
            // repaired parse: only return if not ambiguous
            return potentialMatches.length === 1
              ? { partial: potentialMatches[0] as CHOICE }
              : undefined;
          }
        }
      }
    },

    createElementStreamTransform() {
      return undefined;
    },
  };
};

/**
 * Output specification for unstructured JSON generation.
 * When the model generates a text response, it will return a JSON object.
 *
 * @param name - Optional name of the output that should be generated. Used by some providers for additional LLM guidance, e.g. via tool or schema name.
 * @param description - Optional description of the output that should be generated. Used by some providers for additional LLM guidance, e.g. via tool or schema description.
 *
 * @returns An output specification for generating JSON.
 */
export const json = ({
  name,
  description,
}: {
  /**
   * Optional name of the output that should be generated.
   * Used by some providers for additional LLM guidance, e.g. via tool or schema name.
   */
  name?: string;
  /**
   * Optional description of the output that should be generated.
   * Used by some providers for additional LLM guidance, e.g. via tool or schema description.
   */
  description?: string;
} = {}): Output<JSONValue, JSONValue, never> => {
  return {
    name: 'json',

    responseFormat: Promise.resolve({
      type: 'json' as const,
      ...(name != null && { name }),
      ...(description != null && { description }),
    }),

    async parseCompleteOutput(
      { text }: { text: string },
      context: {
        response: LanguageModelResponseMetadata;
        usage: LanguageModelUsage;
        finishReason: FinishReason;
      },
    ) {
      const parseResult = await safeParseJSON({ text });

      if (!parseResult.success) {
        throw new NoObjectGeneratedError({
          message: 'No object generated: could not parse the response.',
          cause: parseResult.error,
          text,
          response: context.response,
          usage: context.usage,
          finishReason: context.finishReason,
        });
      }

      return parseResult.value;
    },

    async parsePartialOutput({ text }: { text: string }) {
      const result = await parsePartialJson(text);

      switch (result.state) {
        case 'failed-parse':
        case 'undefined-input': {
          return undefined;
        }

        case 'repaired-parse':
        case 'successful-parse': {
          return result.value === undefined
            ? undefined
            : { partial: result.value };
        }
      }
    },

    createElementStreamTransform() {
      return undefined;
    },
  };
};