// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

// GenAI Objective-C API
//
// This is a wrapper around the C++ API, and provides for a set of Objective-C classes with automatic resource management

/* A simple end to end example of how to generate an answer from a prompt:
 *
 * OGAModel* model = [[OGAModel alloc] initWithPath:path error:&error];
 * OGATokenizer* tokenizer = [[OGATokenizer alloc] initWithModel:model error:&error];
 *
 * OGASequences* sequences = [tokenizer encode:@"A great recipe for Kung Pao chicken is " error:&error];
 *
 * OGAGeneratorParams* params = [[OGAGeneratorParams alloc] initWithModel:model error:&error];
 * [params setInputSequences:sequences error:&error];
 * [params setSearchOption:@"max_length" doubleValue:200 error:&error];
 *
 * OGASequences* output_sequences = [model generate:params error:&error];
 * NSString* out_string = [tokenizer decode:[output_sequences sequenceDataAtIndex:0] length:[output_sequences sequenceCountAtIndex:0] error:&error];
 *
 */

NS_ASSUME_NONNULL_BEGIN

@class OGATensor;
@class OGASequences;
@class OGANamedTensors;
@class OGAGeneratorParams;
@class OGATokenizerStream;
@class OGAMultiModalProcessor;

typedef NS_ENUM(NSInteger, OGAElementType) {
  OGAElementTypeUndefined,
  OGAElementTypeFloat32,  // maps to c type float
  OGAElementTypeUint8,    // maps to c type uint8_t
  OGAElementTypeInt8,     // maps to c type int8_t
  OGAElementTypeUint16,   // maps to c type uint16_t
  OGAElementTypeInt16,    // maps to c type int16_t
  OGAElementTypeInt32,    // maps to c type int32_t
  OGAElementTypeInt64,    // maps to c type int64_t
  OGAElementTypeString,   // string type (not currently supported by Oga)
  OGAElementTypeBool,     // maps to c type bool
  OGAElementTypeFloat16,  // IEEE 752-2008 binary16 format, 1 sign bit, 5 bit exponent, 10 bit fraction
  OGAElementTypeFloat64,  // maps to c type double
  OGAElementTypeUint32,   // maps to c type uint32_t
  OGAElementTypeUint64,   // maps to c type uint64_t
};

/** An immutable snapshot of speculative decoding statistics. */
@interface OGASpeculativeStats : NSObject

- (instancetype)init NS_UNAVAILABLE;
- (uint64_t)getCount:(NSString*)name error:(NSError**)error;
- (double)getNumber:(NSString*)name error:(NSError**)error;
- (BOOL)getBool:(NSString*)name error:(NSError**)error;

@end

/**
 * An ORT GenAI config.
 */
@interface OGAConfig : NSObject

- (instancetype)init NS_UNAVAILABLE;

/**
 * Creates a config.
 *
 * @param path The path to the ONNX GenAI model folder.
 * @return The instance, or nil if an error occurs.
 */
- (nullable instancetype)initWithPath:(NSString*)path
                                error:(NSError**)error NS_DESIGNATED_INITIALIZER;

/**
 * Clear the list of providers in the given config
 *
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)clearProvidersWithError:(NSError**)error;

/**
 * Add the provider at the end of the list of providers in the given config if it doesn't already exist.
 * If it already exists, do nothing.
 *
 * @param provider The provider to set on the config.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)appendProvider:(NSString*)provider
                 error:(NSError**)error;

/**
 * Set a provider option.
 *
 * @param provider The provider to set the option on
 * @param key The key of the option to set
 * @param value The value of the option to set
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)setProviderOption:(NSString*)provider
                      key:(NSString*)key
                    value:(NSString*)value
                    error:(NSError**)error;

/**
 * Overlay JSON on top of config file
 *
 * @param json The JSON to overlay on the config.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)overlay:(NSString*)json
          error:(NSError**)error;

@end

/**
 * An ORT GenAI model.
 */
@interface OGAModel : NSObject

- (instancetype)init NS_UNAVAILABLE;

/**
 * Creates a model.
 *
 * @param path The path to the ONNX GenAI model folder.
 * @return The instance, or nil if an error occurs.
 */
- (nullable instancetype)initWithPath:(NSString*)path
                                error:(NSError**)error NS_DESIGNATED_INITIALIZER;

/**
 * Creates a model.
 *
 * @param config The OGAConfig object
 * @return The instance, or nil if an error occurs.
 */
- (nullable instancetype)initWithConfig:(OGAConfig*)config
                                  error:(NSError**)error NS_DESIGNATED_INITIALIZER;
@end

/**
 * An ORT GenAI tokenizer.
 */
@interface OGATokenizer : NSObject
- (instancetype)init NS_UNAVAILABLE;

/**
 * Creates a tokenizer.
 *
 * @param model The model to use.
 * @param error Optional error information set if an error occurs.
 * @return The instance, or nil if an error occurs.
 */
- (nullable instancetype)initWithModel:(OGAModel*)model
                                 error:(NSError**)error NS_DESIGNATED_INITIALIZER;

/**
 * Return the int representation of the BOS token.
 *
 * @return The BOS token id
 */
- (int32_t)getBosTokenId:(NSError**)error;

/**
 * Return the int representations of the array of EOS tokens.
 *
 * @return The array of EOS token ids
 */
- (nullable NSArray<NSNumber*>*)getEosTokenIds:(NSError**)error;

/**
 * Return the int representation of the PAD token.
 *
 * @return The PAD token id
 */
- (int32_t)getPadTokenId:(NSError**)error;

/**
 * Return the BOT (beginning of tool call) token ID. Returns an error if not defined.
 */
- (int32_t)getBotTokenId:(NSError**)error;

/**
 * Return the EOT (end of tool call) token ID. Returns an error if not defined.
 */
- (int32_t)getEotTokenId:(NSError**)error;

/**
 * Return the BOR (beginning of reasoning) token ID. Returns an error if not defined.
 */
- (int32_t)getBorTokenId:(NSError**)error;

/**
 * Return the EOR (end of reasoning) token ID. Returns an error if not defined.
 */
- (int32_t)getEorTokenId:(NSError**)error;

/**
 * Encode text to sequences
 *
 * @param str The text to be encoded.
 * @param error Optional error information set if an error occurs.
 * @return The encoding result, or nil if an error occurs.
 */
- (nullable OGASequences*)encode:(NSString*)str
                           error:(NSError**)error;

/**
 * Decode sequences to text
 *
 * @param tokensData The sequences data to be encoded.
 * @param tokensLength The length of the sequences data to be encoded.
 * @param error Optional error information set if an error occurs.
 * @return The decoding result, or nil if an error occurs.
 */
- (nullable NSString*)decode:(const int32_t*)tokensData
                      length:(size_t)tokensLength
                       error:(NSError**)error;

@end

/**
 * A tokenizer stream is used to decode token strings incrementally, one token at a time.
 */
@interface OGATokenizerStream : NSObject
- (instancetype)init NS_UNAVAILABLE;

/**
 * Creates a tokenizer stream with an underlying tokenizer.
 *
 * @param tokenizer The underlying tokenizer to use.
 * @param error Optional error information set if an error occurs.
 * @return The instance, or nil if an error occurs.
 */
- (nullable instancetype)initWithTokenizer:(OGATokenizer*)tokenizer
                                     error:(NSError**)error NS_DESIGNATED_INITIALIZER;

/**
 * Creates a tokenizer stream with a multi modal processor
 *
 * @param processor The underlying processor to use.
 * @param error Optional error information set if an error occurs.
 * @return The instance, or nil if an error occurs.
 */
- (nullable instancetype)initWithMultiModalProcessor:(OGAMultiModalProcessor*)processor
                                               error:(NSError**)error NS_DESIGNATED_INITIALIZER;
/**
 * Decode one token.
 *
 * @param token The token to be decoded.
 * @param error Optional error information set if an error occurs.
 * @return The decoding result, or nil if an error occurs.
 */
- (nullable NSString*)decode:(int32_t)token
                       error:(NSError**)error;
@end

/**
 * A series of generated sequences
 */
@interface OGASequences : NSObject

- (instancetype)init NS_UNAVAILABLE;

/**
 * The count of generated sequences
 * @param error Optional error information set if an error occurs.
 * @return The count of sequences, or size_t(-1) if an error occurs.
 */
- (size_t)getCountWithError:(NSError**)error NS_SWIFT_NAME(count());

/**
 * Retrieve the sequence data at the given index.
 * @param index The index needed.
 * @param error Optional error information set if an error occurs.
 * @return The sequence data at the given index, or nil if an error occurs.
 */
- (nullable const int32_t*)sequenceDataAtIndex:(size_t)index
                                         error:(NSError**)error;

/**
 * Retrieve the sequence count at the given index.
 * @param index The index needed.
 * @param error Optional error information set if an error occurs.
 * @return The sequence count at the given index, or size_t(-1) if an error occurs.
 */
- (size_t)sequenceCountAtIndex:(size_t)index
                         error:(NSError**)error;

@end

/**
 * The parameters for generation.
 */
@interface OGAGeneratorParams : NSObject
- (instancetype)init NS_UNAVAILABLE;

/**
 * Creates a GeneratorParams from the given model.
 * @param model The model to use for generation.
 * @param error Optional error information set if an error occurs.
 * @return The instance, or nil if an error occurs.
 */
- (nullable instancetype)initWithModel:(OGAModel*)model
                                 error:(NSError**)error NS_DESIGNATED_INITIALIZER;

/**
 * Set double option value.
 * @param key The option key.
 * @param value The option value.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)setSearchOption:(NSString*)key
            doubleValue:(double)value
                  error:(NSError**)error;
/**
 * Set boolean option value.
 * @param key The option key.
 * @param value The option value.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)setSearchOption:(NSString*)key
              boolValue:(BOOL)value
                  error:(NSError**)error;

/**
 * Get numerical value of option.
 * @param key The option key.
 * @param error Optional error information set if an error occurs.
 * @return The option value.
 */
- (double)getSearchNumber:(NSString*)key
                    error:(NSError**)error;
/**
 * Get boolean value of option.
 * @param key The option key.
 * @param error Optional error information set if an error occurs.
 * @return The option value.
 */
- (BOOL)getSearchBool:(NSString*)key
                error:(NSError**)error;

/**
 * Set a numerical speculative decoding option.
 * @param key The option key.
 * @param value The option value.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)setSpeculativeNumber:(NSString*)key
                 doubleValue:(double)value
                       error:(NSError**)error;

/**
 * Get a numerical speculative decoding option.
 * @param key The option key.
 * @param error Optional error information set if an error occurs.
 * @return The option value.
 */
- (double)getSpeculativeNumber:(NSString*)key
                         error:(NSError**)error;
@end

/**
 * The main generator interface that can be used for generation loop.
 */
@interface OGAGenerator : NSObject

- (instancetype)init NS_UNAVAILABLE;
/**
 * Creates a generator.
 *
 * @param model The model to use.
 * @param params The generation params to use.
 * @param error Optional error information set if an error occurs.
 * @return The instance, or nil if an error occurs.
 */
- (nullable instancetype)initWithModel:(OGAModel*)model
                                params:(OGAGeneratorParams*)params
                                 error:(NSError**)error NS_DESIGNATED_INITIALIZER;
/**
 * Whether generation is done.
 * @param error Optional error information set if an error occurs.
 * @return The result, or false if an error occurs.
 */
- (BOOL)isDoneWithError:(NSError**)error __attribute__((swift_error(nonnull_error)));

/**
 * Set input with NamedTensors type.
 * @param namedTensors The named tensors.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)setInputs:(OGANamedTensors*)namedTensors
            error:(NSError**)error;

/**
 * Set input with name and corresponding tensor.
 * @param name The input name.
 * @param tensor The tensor.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)setModelInput:(NSString*)name
               tensor:(OGATensor*)tensor
                error:(NSError**)error;

/**
 * Appends token sequences to the generator.
 * @param sequences The sequences to append.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)appendTokenSequences:(OGASequences*)sequences error:(NSError**)error;

/**
 * Appends token sequences to the generator.
 * @param tokens The tokens to append.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)appendTokens:(NSArray<NSNumber*>*)tokens error:(NSError**)error;

/**
 * Get the number of tokens in the generator.
 * @param error Optional error information set if an error occurs.
 * @return The number of tokens in the generator
 */
- (size_t)tokenCount:(NSError**)error;

/**
 * Rewinds the generator to the given length.
 * @param newLength The desired length in tokens after rewinding.
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)rewindTo:(size_t)newLength error:(NSError**)error;

/**
 * Generate next token
 * @param error Optional error information set if an error occurs.
 */
- (BOOL)generateNextTokenWithError:(NSError**)error;

/**
 * Get the input tensor.
 * @param name The input name.
 * @param error Optional error information set if an error occurs.
 * @return The result, or nil if an error occurs.
 */
- (nullable OGATensor*)getInput:(NSString*)name error:(NSError**)error;

/**
 * Get the output tensor.
 * @param name The output name.
 * @param error Optional error information set if an error occurs.
 * @return The result, or nil if an error occurs.
 */
- (nullable OGATensor*)getOutput:(NSString*)name error:(NSError**)error;

/**
 * Retrieve the sequence data at the given index.
 * @param index The index needed.
 * @param error Optional error information set if an error occurs.
 * @return The sequence data at the given index, or nil if an error occurs.
 */
- (nullable const int32_t*)sequenceDataAtIndex:(size_t)index
                                         error:(NSError**)error;

/**
 * Retrieve the sequence count at the given index.
 * @param index The index needed.
 * @param error Optional error information set if an error occurs.
 * @return The sequence count at the given index, or size_t(-1) if an error occurs.
 */
- (size_t)sequenceCountAtIndex:(size_t)index
                         error:(NSError**)error;

/** Get an immutable snapshot of the accumulated speculative decoding statistics. */
- (nullable OGASpeculativeStats*)getSpeculativeStatsWithError:(NSError**)error;

/**
 * Clean up the resource before process exits.
 */
+ (void)shutdown;

/**
 * Enable or disable non-essential ONNX Runtime GenAI telemetry events.
 *
 * @param enabled Whether telemetry events are enabled.
 */
+ (void)setTelemetryEnabled:(BOOL)enabled;

@end

@interface OGATensor : NSObject

- (instancetype)init NS_UNAVAILABLE;
- (nullable instancetype)initWithDataPointer:(void*)data
                                       shape:(NSArray<NSNumber*>*)shape
                                        type:(OGAElementType)elementType
                                       error:(NSError**)error;
- (OGAElementType)getTypeWithError:(NSError**)error NS_SWIFT_NAME(type());
- (nullable void*)getDataPointerWithError:(NSError**)error NS_SWIFT_NAME(dataPointer());

@end

@interface OGANamedTensors : NSObject

- (instancetype)init NS_UNAVAILABLE;

@end

@interface OGAImages : NSObject

- (instancetype)init NS_UNAVAILABLE;
- (nullable instancetype)initWithPath:(NSArray<NSString*>*)paths
                                error:(NSError**)error NS_DESIGNATED_INITIALIZER;

@end

@interface OGAAudios : NSObject

- (instancetype)init NS_UNAVAILABLE;
- (nullable instancetype)initWithPath:(NSArray<NSString*>*)paths
                                error:(NSError**)error NS_DESIGNATED_INITIALIZER;

@end

@interface OGAMultiModalProcessor : NSObject

- (instancetype)init NS_UNAVAILABLE;
- (nullable instancetype)initWithModel:(OGAModel*)model
                                 error:(NSError**)error NS_DESIGNATED_INITIALIZER;

- (nullable OGANamedTensors*)processImages:(NSString*)prompt
                                    images:(OGAImages*)images
                                     error:(NSError**)error;

- (nullable OGANamedTensors*)processImages:(NSArray<NSString*>*)prompts
                                    images:(OGAImages*)images
                                     error:(NSError**)error;

- (nullable OGANamedTensors*)processAudios:(NSString*)prompt
                                    audios:(OGAAudios*)audios
                                     error:(NSError**)error;

- (nullable OGANamedTensors*)processAudios:(NSArray<NSString*>*)prompts
                                    audios:(OGAAudios*)audios
                                     error:(NSError**)error;

- (nullable OGANamedTensors*)processImagesAndAudios:(NSString*)prompt
                                             images:(OGAImages*)images
                                             audios:(OGAAudios*)audios
                                              error:(NSError**)error;

- (nullable OGANamedTensors*)processImagesAndAudios:(NSArray<NSString*>*)prompts
                                             images:(OGAImages*)images
                                             audios:(OGAAudios*)audios
                                              error:(NSError**)error;

- (nullable NSString*)decode:(const int32_t*)tokensData
                      length:(size_t)tokensLength
                       error:(NSError**)error;
@end

NS_ASSUME_NONNULL_END
