import { NativeEventEmitter, DeviceEventEmitter, Platform } from 'react-native';
import RNLlama from './NativeRNLlama';
import { SchemaGrammarConverter, convertJsonSchemaToGrammar } from './grammar';
import { formatChat } from './chat';
export { SchemaGrammarConverter, convertJsonSchemaToGrammar };
const EVENT_ON_TOKEN = '@RNLlama_onToken';
let EventEmitter;
if (Platform.OS === 'ios') {
  // @ts-ignore
  EventEmitter = new NativeEventEmitter(RNLlama);
}
if (Platform.OS === 'android') {
  EventEmitter = DeviceEventEmitter;
}
export class LlamaContext {
  gpu = false;
  reasonNoGPU = '';
  model = {};
  constructor(_ref) {
    let {
      contextId,
      gpu,
      reasonNoGPU,
      model
    } = _ref;
    this.id = contextId;
    this.gpu = gpu;
    this.reasonNoGPU = reasonNoGPU;
    this.model = model;
  }

  /**
   * Load cached prompt & completion state from a file.
   */
  async loadSession(filepath) {
    let path = filepath;
    if (path.startsWith('file://')) path = path.slice(7);
    return RNLlama.loadSession(this.id, path);
  }

  /**
   * Save current cached prompt & completion state to a file.
   */
  async saveSession(filepath, options) {
    return RNLlama.saveSession(this.id, filepath, (options === null || options === void 0 ? void 0 : options.tokenSize) || -1);
  }
  async getFormattedChat(messages) {
    var _this$model;
    const chat = formatChat(messages);
    return RNLlama.getFormattedChat(this.id, chat, (_this$model = this.model) !== null && _this$model !== void 0 && _this$model.isChatTemplateSupported ? undefined : 'chatml');
  }
  async completion(params, callback) {
    let finalPrompt = params.prompt;
    if (params.messages) {
      // messages always win
      finalPrompt = await this.getFormattedChat(params.messages);
    }
    let tokenListener = callback && EventEmitter.addListener(EVENT_ON_TOKEN, evt => {
      const {
        contextId,
        tokenResult
      } = evt;
      if (contextId !== this.id) return;
      callback(tokenResult);
    });
    if (!finalPrompt) throw new Error('Prompt is required');
    const promise = RNLlama.completion(this.id, {
      ...params,
      prompt: finalPrompt,
      emit_partial_completion: !!callback
    });
    return promise.then(completionResult => {
      var _tokenListener;
      (_tokenListener = tokenListener) === null || _tokenListener === void 0 ? void 0 : _tokenListener.remove();
      tokenListener = null;
      return completionResult;
    }).catch(err => {
      var _tokenListener2;
      (_tokenListener2 = tokenListener) === null || _tokenListener2 === void 0 ? void 0 : _tokenListener2.remove();
      tokenListener = null;
      throw err;
    });
  }
  stopCompletion() {
    return RNLlama.stopCompletion(this.id);
  }
  tokenize(text) {
    return RNLlama.tokenize(this.id, text);
  }
  detokenize(tokens) {
    return RNLlama.detokenize(this.id, tokens);
  }
  embedding(text) {
    return RNLlama.embedding(this.id, text);
  }
  async bench(pp, tg, pl, nr) {
    const result = await RNLlama.bench(this.id, pp, tg, pl, nr);
    const [modelDesc, modelSize, modelNParams, ppAvg, ppStd, tgAvg, tgStd] = JSON.parse(result);
    return {
      modelDesc,
      modelSize,
      modelNParams,
      ppAvg,
      ppStd,
      tgAvg,
      tgStd
    };
  }
  async release() {
    return RNLlama.releaseContext(this.id);
  }
}
export async function setContextLimit(limit) {
  return RNLlama.setContextLimit(limit);
}
export async function initLlama(_ref2) {
  let {
    model,
    is_model_asset: isModelAsset,
    ...rest
  } = _ref2;
  let path = model;
  if (path.startsWith('file://')) path = path.slice(7);
  const {
    contextId,
    gpu,
    reasonNoGPU,
    model: modelDetails
  } = await RNLlama.initContext({
    model: path,
    is_model_asset: !!isModelAsset,
    ...rest
  });
  return new LlamaContext({
    contextId,
    gpu,
    reasonNoGPU,
    model: modelDetails
  });
}
export async function releaseAllLlama() {
  return RNLlama.releaseAllContexts();
}
//# sourceMappingURL=index.js.map