import { Platform } from 'react-native'
import RNLlama from './NativeRNLlama'
import './jsi'
import type {
  NativeContextParams,
  NativeLlamaContext,
  NativeCompletionParams,
  NativeParallelCompletionParams,
  NativeCompletionTokenProb,
  NativeCompletionResult,
  NativeTokenizeResult,
  NativeEmbeddingResult,
  NativeSessionLoadResult,
  NativeEmbeddingParams,
  NativeRerankParams,
  NativeRerankResult,
  NativeCompletionTokenProbItem,
  NativeCompletionResultTimings,
  JinjaFormattedChatResult,
  FormattedChatResult,
  NativeImageProcessingResult,
  NativeLlamaChatMessage,
  NativeBackendDeviceInfo,
  NativeSpeculativeConfig,
  NativeSpeculativeParams,
  NativeSpeculativeType,
  ParallelStatus,
  ParallelRequestStatus,
} from './types'
import { BUILD_NUMBER, BUILD_COMMIT } from './version'
import type { SpeakerPayload } from './tts-voices'

export type RNLlamaMessagePart = {
  type: string
  text?: string
  image_url?: {
    url?: string
  }
  input_audio?: {
    format: string
    data?: string
    url?: string
  }
}

export type RNLlamaOAICompatibleMessage = {
  role: string
  content?: string | RNLlamaMessagePart[]
  reasoning_content?: string
}

export type {
  NativeContextParams,
  NativeLlamaContext,
  NativeCompletionParams,
  NativeParallelCompletionParams,
  NativeCompletionTokenProb,
  NativeCompletionResult,
  NativeTokenizeResult,
  NativeEmbeddingResult,
  NativeSessionLoadResult,
  NativeEmbeddingParams,
  NativeRerankParams,
  NativeRerankResult,
  NativeCompletionTokenProbItem,
  NativeCompletionResultTimings,
  FormattedChatResult,
  JinjaFormattedChatResult,
  NativeImageProcessingResult,
  NativeBackendDeviceInfo,
  NativeSpeculativeConfig,
  NativeSpeculativeParams,
  NativeSpeculativeType,
  ParallelStatus,
  ParallelRequestStatus,
}

export const RNLLAMA_MTMD_DEFAULT_MEDIA_MARKER = '<__media__>'

export type { TTSCapabilities } from './tts'
export {
  lookupVoice as getTTSVoice,
  listVoices as listTTSVoices,
  listLanguages as listTTSLanguages,
} from './tts-voices'
export type {
  OuteTTSWord,
  OuteTTSSpeaker,
  NeuTTSSpeaker,
  SpeakerPayload,
} from './tts-voices'

const logListeners: Array<(level: string, text: string) => void> = []
const emitNativeLog = (level: string, text: string) => {
  logListeners.forEach((listener) => listener(level, text))
}

const jsiBindingKeys = [
  'llamaInitContext',
  'llamaReleaseContext',
  'llamaReleaseAllContexts',
  'llamaModelInfo',
  'llamaGetBackendDevicesInfo',
  'llamaLoadSession',
  'llamaSaveSession',
  'llamaTokenize',
  'llamaDetokenize',
  'llamaGetFormattedChat',
  'llamaEmbedding',
  'llamaRerank',
  'llamaBench',
  'llamaToggleNativeLog',
  'llamaSetContextLimit',
  'llamaCompletion',
  'llamaStopCompletion',
  'llamaApplyLoraAdapters',
  'llamaRemoveLoraAdapters',
  'llamaGetLoadedLoraAdapters',
  'llamaInitMultimodal',
  'llamaIsMultimodalEnabled',
  'llamaGetMultimodalSupport',
  'llamaReleaseMultimodal',
  'llamaInitVocoder',
  'llamaIsVocoderEnabled',
  'llamaGetFormattedAudioCompletion',
  'llamaGetTTSCapabilities',
  'llamaDecodeAudioTokens',
  'llamaGenerateAudioCodes',
  'llamaCreateSpeaker',
  'llamaBakeSpeaker',
  'llamaReleaseSpeaker',
  'llamaDecodeAudioEmbeddings',
  'llamaGetAudioSampleRate',
  'llamaReleaseVocoder',
  'llamaClearCache',
  'llamaEnableParallelMode',
  'llamaQueueCompletion',
  'llamaCancelRequest',
  'llamaQueueEmbedding',
  'llamaQueueRerank',
  'llamaGetParallelStatus',
  'llamaSubscribeParallelStatus',
  'llamaUnsubscribeParallelStatus',
] as const

type JsiBindingKey = (typeof jsiBindingKeys)[number]
type JsiBindings = { [K in JsiBindingKey]: NonNullable<(typeof globalThis)[K]> }

let jsiBindings: JsiBindings | null = null

const bindJsiFromGlobal = () => {
  const bindings: Partial<JsiBindings> = {}
  const missing: string[] = []

  jsiBindingKeys.forEach((key) => {
    const value = global[key]
    if (typeof value === 'function') {
      ;(bindings as Record<string, unknown>)[key] =
        value as JsiBindings[typeof key]
      delete global[key]
    } else {
      missing.push(key)
    }
  })

  if (missing.length > 0) {
    throw new Error(`[RNLlama] Missing JSI bindings: ${missing.join(', ')}`)
  }

  jsiBindings = bindings as JsiBindings
}

const getJsi = (): JsiBindings => {
  if (!jsiBindings) {
    throw new Error('JSI bindings not installed')
  }
  return jsiBindings
}

// JSI Installation
let isJsiInstalled = false
export const installJsi = async () => {
  if (isJsiInstalled) return
  if (typeof global.llamaInitContext !== 'function') {
    const installed = await RNLlama.install()
    if (!installed && typeof global.llamaInitContext !== 'function') {
      throw new Error('JSI bindings not installed')
    }
  }
  bindJsiFromGlobal()
  isJsiInstalled = true
}

export type ToolCall = {
  type: 'function'
  id?: string
  function: {
    name: string
    arguments: string // JSON string
  }
}

export type TokenData = {
  token: string
  completion_probabilities?: Array<NativeCompletionTokenProb>
  // Parsed content from accumulated text
  content?: string
  reasoning_content?: string
  tool_calls?: Array<ToolCall>
  accumulated_text?: string
  requestId?: number
}

export type ContextParams = Omit<
  NativeContextParams,
  'flash_attn_type' | 'cache_type_k' | 'cache_type_v' | 'pooling_type'
> & {
  flash_attn_type?: 'auto' | 'on' | 'off'
  cache_type_k?:
    | 'f16'
    | 'f32'
    | 'q8_0'
    | 'q4_0'
    | 'q4_1'
    | 'iq4_nl'
    | 'q5_0'
    | 'q5_1'
  cache_type_v?:
    | 'f16'
    | 'f32'
    | 'q8_0'
    | 'q4_0'
    | 'q4_1'
    | 'iq4_nl'
    | 'q5_0'
    | 'q5_1'
  pooling_type?: 'none' | 'mean' | 'cls' | 'last' | 'rank'
}

const validCacheTypes = [
  'f16',
  'f32',
  'bf16',
  'q8_0',
  'q4_0',
  'q4_1',
  'iq4_nl',
  'q5_0',
  'q5_1',
]

export type EmbeddingParams = NativeEmbeddingParams

export type RerankParams = {
  normalize?: number
}

export type RerankResult = {
  score: number
  index: number
  document?: string
}

export type CompletionResponseFormat = {
  type: 'text' | 'json_object' | 'json_schema'
  json_schema?: {
    strict?: boolean
    schema: object
  }
  schema?: object // for json_object type
}

export type ChatTemplateKwargs = Record<string, string | number | boolean>

export type CompletionBaseParams = {
  prompt?: string
  messages?: RNLlamaOAICompatibleMessage[]
  chatTemplate?: string // deprecated
  chat_template?: string
  jinja?: boolean
  tools?: object
  parallel_tool_calls?: object
  tool_choice?: string
  response_format?: CompletionResponseFormat
  media_paths?: string | string[]
  add_generation_prompt?: boolean
  /*
   * Timestamp in seconds since epoch to apply to chat template's strftime_now
   */
  now?: string | number
  chat_template_kwargs?: ChatTemplateKwargs
  /**
   * When enabled, forces the chat parser to treat the entire model output as
   * plain content, skipping separate parsing of reasoning tokens and tool calls.
   * Also bypasses jinja template validation so templates that only accept typed
   * content (e.g. TranslateGemma) are not rejected during capability detection.
   */
  force_pure_content?: boolean
  /**
   * Prefill text to be used for chat parsing (Generation Prompt + Content)
   * Used for if last assistant message is for prefill purpose
   */
  prefill_text?: string
}

export type CompletionParams = Omit<
  NativeCompletionParams,
  'emit_partial_completion' | 'prompt'
> &
  CompletionBaseParams

/**
 * Parameters for parallel completion requests.
 * Extends CompletionParams with parallel-mode specific options like state management.
 */
export type ParallelCompletionParams = Omit<
  NativeParallelCompletionParams,
  'emit_partial_completion' | 'prompt'
> &
  CompletionBaseParams

type ReasoningBudgetChatMetadata = {
  thinking_start_tag?: string
  thinking_end_tag?: string
}

type NativeCompletionRequestParams = NativeCompletionParams &
  ReasoningBudgetChatMetadata

type NativeParallelCompletionRequestParams = NativeParallelCompletionParams &
  ReasoningBudgetChatMetadata

export type BenchResult = {
  nKvMax: number
  nBatch: number
  nUBatch: number
  flashAttn: number
  isPpShared: number
  nGpuLayers: number
  nThreads: number
  nThreadsBatch: number
  pp: number
  tg: number
  pl: number
  nKv: number
  tPp: number
  speedPp: number
  tTg: number
  speedTg: number
  t: number
  speed: number
}

const getJsonSchema = (responseFormat?: CompletionResponseFormat) => {
  if (responseFormat?.type === 'json_schema') {
    return responseFormat.json_schema?.schema
  }
  if (responseFormat?.type === 'json_object') {
    return responseFormat.schema || {}
  }
  return null
}

export class LlamaSpeaker {
  readonly id: number

  readonly family: string

  rows: number

  baked: boolean

  private ctxId: number

  constructor(ctxId: number, h: { id: number; family: string; rows: number; baked: boolean }) {
    this.ctxId = ctxId
    this.id = h.id
    this.family = h.family
    this.rows = h.rows
    this.baked = h.baked
  }

  async bake(): Promise<void> {
    const { llamaBakeSpeaker } = getJsi()
    const r = await llamaBakeSpeaker(this.ctxId, this.id)
    this.rows = r.rows
    this.baked = r.baked
  }

  async release(): Promise<void> {
    const { llamaReleaseSpeaker } = getJsi()
    await llamaReleaseSpeaker(this.ctxId, this.id)
  }
}

export class LlamaContext {
  id: number

  gpu: boolean = false

  reasonNoGPU: string = ''

  devices: NativeLlamaContext['devices']

  model: NativeLlamaContext['model']

  androidLib: NativeLlamaContext['androidLib']

  systemInfo: NativeLlamaContext['systemInfo']

  /**
   * Parallel processing namespace for non-blocking queue operations
   */
  parallel = {
    /**
     * Queue a completion request for parallel processing (non-blocking)
     * @param params Parallel completion parameters (includes state management)
     * @param onToken Callback fired for each generated token
     * @returns Promise resolving to object with requestId, promise (resolves to completion result), and stop function
     */
    completion: async (
      params: ParallelCompletionParams,
      onToken?: (requestId: number, data: TokenData) => void,
    ): Promise<{
      requestId: number
      promise: Promise<NativeCompletionResult>
      stop: () => Promise<void>
    }> => {
      const { llamaQueueCompletion, llamaCancelRequest } = getJsi()
      const nativeParams: NativeParallelCompletionRequestParams = {
        ...params,
        prompt: params.prompt || '',
        emit_partial_completion: true, // Always emit for queued requests
      }

      // Process messages same as completion()
      if (params.messages) {
        const formattedResult = await this.getFormattedChat(
          params.messages,
          params.chat_template || params.chatTemplate,
          {
            jinja: params.jinja,
            tools: params.tools,
            parallel_tool_calls: params.parallel_tool_calls,
            tool_choice: params.tool_choice,
            response_format: params.response_format,
            enable_thinking: params.enable_thinking,
            reasoning_format: params.reasoning_format,
            add_generation_prompt: params.add_generation_prompt,
            now: params.now,
            chat_template_kwargs: params.chat_template_kwargs,
            force_pure_content: params.force_pure_content,
          },
        )
        if (formattedResult.type === 'jinja') {
          const jinjaResult = formattedResult as JinjaFormattedChatResult
          nativeParams.prompt = jinjaResult.prompt || ''
          if (typeof jinjaResult.chat_format === 'number')
            nativeParams.chat_format = jinjaResult.chat_format
          if (jinjaResult.grammar) nativeParams.grammar = jinjaResult.grammar
          if (typeof jinjaResult.grammar_lazy === 'boolean')
            nativeParams.grammar_lazy = jinjaResult.grammar_lazy
          if (jinjaResult.grammar_triggers)
            nativeParams.grammar_triggers = jinjaResult.grammar_triggers
          if (jinjaResult.preserved_tokens)
            nativeParams.preserved_tokens = jinjaResult.preserved_tokens
          if (jinjaResult.additional_stops) {
            if (!nativeParams.stop) nativeParams.stop = []
            nativeParams.stop.push(...jinjaResult.additional_stops)
          }
          if (jinjaResult.has_media) {
            nativeParams.media_paths = jinjaResult.media_paths
          }
          if (typeof jinjaResult.generation_prompt === 'string')
            nativeParams.generation_prompt = jinjaResult.generation_prompt
          if (typeof jinjaResult.thinking_forced_open === 'boolean')
            nativeParams.thinking_forced_open = jinjaResult.thinking_forced_open
          if (typeof jinjaResult.thinking_start_tag === 'string')
            nativeParams.thinking_start_tag = jinjaResult.thinking_start_tag
          if (typeof jinjaResult.thinking_end_tag === 'string')
            nativeParams.thinking_end_tag = jinjaResult.thinking_end_tag
          if (jinjaResult.chat_parser)
            nativeParams.chat_parser = jinjaResult.chat_parser
        } else if (formattedResult.type === 'llama-chat') {
          const llamaChatResult = formattedResult as FormattedChatResult
          nativeParams.prompt = llamaChatResult.prompt || ''
          if (llamaChatResult.has_media) {
            nativeParams.media_paths = llamaChatResult.media_paths
          }
        }
      } else {
        nativeParams.prompt = params.prompt || ''
      }

      if (!nativeParams.media_paths && params.media_paths) {
        nativeParams.media_paths = params.media_paths
      }

      if (params.response_format && !nativeParams.grammar) {
        const jsonSchema = getJsonSchema(params.response_format)
        if (jsonSchema) nativeParams.json_schema = JSON.stringify(jsonSchema)
      }

      // Empty prompt is valid for embedding-injection flows (see completion()).
      if (!nativeParams.prompt && !params.embedding && !nativeParams.media_paths)
        throw new Error('Prompt is required')

      return new Promise(async (resolveOuter, rejectOuter) => {
        try {
          let resolveResult: (
            value: NativeCompletionResult | PromiseLike<NativeCompletionResult>,
          ) => void
          let rejectResult: (reason?: any) => void

          const resultPromise = new Promise<NativeCompletionResult>(
            (res, rej) => {
              resolveResult = res
              rejectResult = rej
            },
          )

          const { requestId } = await llamaQueueCompletion(
            this.id,
            nativeParams,
            (tokenResult, reqId) => {
              if (onToken) onToken(reqId, tokenResult)
            },
            (result) => {
              if (result.error) {
                rejectResult(new Error(result.error))
              } else {
                resolveResult(result)
              }
            },
          )

          resolveOuter({
            requestId,
            promise: resultPromise,
            stop: async () => {
              await llamaCancelRequest(this.id, requestId)
            },
          })
        } catch (e) {
          rejectOuter(e)
        }
      })
    },

    /**
     * Queue an embedding request for parallel processing (non-blocking)
     * @param text Text to embed
     * @param params Optional embedding parameters
     * @returns Promise resolving to object with requestId and promise (resolves to embedding result)
     */
    embedding: async (
      text: string,
      params?: EmbeddingParams,
    ): Promise<{
      requestId: number
      promise: Promise<NativeEmbeddingResult>
    }> =>
      new Promise(async (resolveOuter, rejectOuter) => {
        const { llamaQueueEmbedding } = getJsi()
        try {
          let resolveResult: (value: NativeEmbeddingResult) => void
          const resultPromise = new Promise<NativeEmbeddingResult>((res) => {
            resolveResult = res
          })

          const { requestId } = await llamaQueueEmbedding(
            this.id,
            text,
            params || {},
            (embedding) => {
              resolveResult({ embedding })
            },
          )

          resolveOuter({
            requestId,
            promise: resultPromise,
          })
        } catch (e) {
          rejectOuter(e)
        }
      }),

    /**
     * Queue rerank requests for parallel processing (non-blocking)
     * @param query The query text to rank documents against
     * @param documents Array of document texts to rank
     * @param params Optional reranking parameters
     * @returns Promise resolving to object with requestId and promise (resolves to rerank results)
     */
    rerank: async (
      query: string,
      documents: string[],
      params?: RerankParams,
    ): Promise<{
      requestId: number
      promise: Promise<RerankResult[]>
    }> =>
      new Promise(async (resolveOuter, rejectOuter) => {
        const { llamaQueueRerank } = getJsi()
        try {
          let resolveResult: (value: RerankResult[]) => void
          const resultPromise = new Promise<RerankResult[]>((res) => {
            resolveResult = res
          })

          const { requestId } = await llamaQueueRerank(
            this.id,
            query,
            documents,
            params || {},
            (results) => {
              const sortedResults = results
                .map((result: NativeRerankResult) => ({
                  ...result,
                  document: documents[result.index],
                }))
                .sort((a: RerankResult, b: RerankResult) => b.score - a.score)
              resolveResult(sortedResults)
            },
          )

          resolveOuter({
            requestId,
            promise: resultPromise,
          })
        } catch (e) {
          rejectOuter(e)
        }
      }),

    enable: (config?: { n_parallel?: number; n_batch?: number }) =>
      getJsi().llamaEnableParallelMode(this.id, { enabled: true, ...config }),

    disable: () =>
      getJsi().llamaEnableParallelMode(this.id, { enabled: false }),

    configure: (config: { n_parallel?: number; n_batch?: number }) =>
      getJsi().llamaEnableParallelMode(this.id, { enabled: true, ...config }),

    /**
     * Get current parallel processing status (one-time snapshot)
     * @returns Promise resolving to current parallel status
     */
    getStatus: async (): Promise<ParallelStatus> => {
      const { llamaGetParallelStatus } = getJsi()
      return llamaGetParallelStatus(this.id)
    },

    /**
     * Subscribe to parallel processing status changes
     * @param callback Called whenever parallel status changes
     * @returns Object with remove() method to unsubscribe
     */
    subscribeToStatus: async (
      callback: (status: ParallelStatus) => void,
    ): Promise<{ remove: () => void }> => {
      const { llamaSubscribeParallelStatus, llamaUnsubscribeParallelStatus } =
        getJsi()
      const { subscriberId } = await llamaSubscribeParallelStatus(
        this.id,
        callback,
      )

      return {
        remove: () => {
          llamaUnsubscribeParallelStatus(this.id, subscriberId)
        },
      }
    },
  }

  constructor({
    contextId,
    gpu,
    devices,
    reasonNoGPU,
    model,
    androidLib,
    systemInfo,
  }: NativeLlamaContext) {
    this.id = contextId
    this.gpu = gpu
    this.devices = devices
    this.reasonNoGPU = reasonNoGPU
    this.model = model
    this.androidLib = androidLib
    this.systemInfo = systemInfo
  }

  async loadSession(filepath: string): Promise<NativeSessionLoadResult> {
    const { llamaLoadSession } = getJsi()
    let path = filepath
    if (path.startsWith('file://')) path = path.slice(7)
    return llamaLoadSession(this.id, path)
  }

  async saveSession(
    filepath: string,
    options?: { tokenSize: number },
  ): Promise<number> {
    const { llamaSaveSession } = getJsi()
    return llamaSaveSession(this.id, filepath, options?.tokenSize || -1)
  }

  isLlamaChatSupported(): boolean {
    return !!this.model.chatTemplates.llamaChat
  }

  isJinjaSupported(): boolean {
    const { jinja } = this.model.chatTemplates
    return !!jinja?.toolUse || !!jinja?.default
  }

  async getFormattedChat(
    messages: RNLlamaOAICompatibleMessage[],
    template?: string | null,
    params?: {
      jinja?: boolean
      response_format?: CompletionResponseFormat
      tools?: object
      parallel_tool_calls?: object
      tool_choice?: string
      enable_thinking?: boolean
      reasoning_format?: 'none' | 'auto' | 'deepseek'
      add_generation_prompt?: boolean
      now?: string | number
      chat_template_kwargs?: ChatTemplateKwargs
      force_pure_content?: boolean
    },
  ): Promise<FormattedChatResult | JinjaFormattedChatResult> {
    const mediaPaths: string[] = []
    const chat = messages.map((msg) => {
      if (Array.isArray(msg.content)) {
        const content = msg.content.map((part) => {
          if (part.type === 'image_url') {
            let path = part.image_url?.url || ''
            if (path?.startsWith('file://')) path = path.slice(7)
            mediaPaths.push(path)
            return {
              type: 'text',
              text: RNLLAMA_MTMD_DEFAULT_MEDIA_MARKER,
            }
          } else if (part.type === 'input_audio') {
            const { input_audio: audio } = part
            if (!audio) throw new Error('input_audio is required')

            const { format } = audio
            if (format != 'wav' && format != 'mp3') {
              throw new Error(`Unsupported audio format: ${format}`)
            }
            if (audio.url) {
              const path = audio.url.replace(/file:\/\//, '')
              mediaPaths.push(path)
            } else if (audio.data) {
              mediaPaths.push(audio.data)
            }
            return {
              type: 'text',
              text: RNLLAMA_MTMD_DEFAULT_MEDIA_MARKER,
            }
          }
          return part
        })

        return {
          ...msg,
          content,
        }
      }
      return msg
    }) as NativeLlamaChatMessage[]

    const forcePureContent = params?.force_pure_content ?? false
    // When force_pure_content is set, accept any model that has a chat_template
    // string in its metadata without requiring template validation to pass.
    const hasChatTemplate = !!(this.model.metadata as Record<string, unknown>)[
      'tokenizer.chat_template'
    ]
    const useJinja =
      (forcePureContent ? hasChatTemplate : this.isJinjaSupported()) &&
      (params?.jinja ?? true)
    let tmpl
    if (template) tmpl = template
    const jsonSchema = getJsonSchema(params?.response_format)

    const { llamaGetFormattedChat } = getJsi()
    const result = await llamaGetFormattedChat(
      this.id,
      JSON.stringify(chat),
      tmpl,
      {
        jinja: useJinja,
        json_schema: jsonSchema ? JSON.stringify(jsonSchema) : undefined,
        tools: params?.tools ? JSON.stringify(params.tools) : undefined,
        parallel_tool_calls: params?.parallel_tool_calls
          ? JSON.stringify(params.parallel_tool_calls)
          : undefined,
        tool_choice: params?.tool_choice,
        enable_thinking: params?.enable_thinking ?? true,
        reasoning_format: params?.reasoning_format ?? 'none',
        add_generation_prompt: params?.add_generation_prompt,
        now:
          typeof params?.now === 'number' ? params.now.toString() : params?.now,
        chat_template_kwargs: params?.chat_template_kwargs
          ? JSON.stringify(
              Object.entries(params.chat_template_kwargs).reduce(
                (acc, [key, value]) => {
                  acc[key] = JSON.stringify(value)
                  return acc
                },
                {} as Record<string, any>,
              ),
            )
          : undefined,
        force_pure_content: forcePureContent,
      },
    )
    if (!useJinja) {
      return {
        type: 'llama-chat',
        prompt: result as string,
        has_media: mediaPaths.length > 0,
        media_paths: mediaPaths,
      }
    }
    const jinjaResult = result as JinjaFormattedChatResult
    jinjaResult.type = 'jinja'
    jinjaResult.has_media = mediaPaths.length > 0
    jinjaResult.media_paths = mediaPaths
    return jinjaResult
  }

  async completion(
    params: CompletionParams & { speaker?: LlamaSpeaker },
    callback?: (data: TokenData) => void,
  ): Promise<NativeCompletionResult> {
    const nativeParams: NativeCompletionRequestParams & {
      speakerId?: number
    } = {
      ...params,
      prompt: params.prompt || '',
      emit_partial_completion: !!callback,
      ...(params.speaker !== undefined
        ? { speakerId: params.speaker.id }
        : {}),
    }

    if (params.messages) {
      const formattedResult = await this.getFormattedChat(
        params.messages,
        params.chat_template || params.chatTemplate,
        {
          jinja: params.jinja,
          tools: params.tools,
          parallel_tool_calls: params.parallel_tool_calls,
          tool_choice: params.tool_choice,
          response_format: params.response_format,
          enable_thinking: params.enable_thinking,
          reasoning_format: params.reasoning_format,
          add_generation_prompt: params.add_generation_prompt,
          now: params.now,
          chat_template_kwargs: params.chat_template_kwargs,
          force_pure_content: params.force_pure_content,
        },
      )
      if (formattedResult.type === 'jinja') {
        const jinjaResult = formattedResult as JinjaFormattedChatResult

        nativeParams.prompt = jinjaResult.prompt || ''
        if (typeof jinjaResult.chat_format === 'number')
          nativeParams.chat_format = jinjaResult.chat_format
        if (jinjaResult.grammar) nativeParams.grammar = jinjaResult.grammar
        if (typeof jinjaResult.grammar_lazy === 'boolean')
          nativeParams.grammar_lazy = jinjaResult.grammar_lazy
        if (jinjaResult.grammar_triggers)
          nativeParams.grammar_triggers = jinjaResult.grammar_triggers
        if (jinjaResult.preserved_tokens)
          nativeParams.preserved_tokens = jinjaResult.preserved_tokens
        if (jinjaResult.additional_stops) {
          if (!nativeParams.stop) nativeParams.stop = []
          nativeParams.stop.push(...jinjaResult.additional_stops)
        }
        if (jinjaResult.has_media) {
          nativeParams.media_paths = jinjaResult.media_paths
        }
        if (typeof jinjaResult.generation_prompt === 'string')
          nativeParams.generation_prompt = jinjaResult.generation_prompt
        if (typeof jinjaResult.thinking_forced_open === 'boolean')
          nativeParams.thinking_forced_open = jinjaResult.thinking_forced_open
        if (typeof jinjaResult.thinking_start_tag === 'string')
          nativeParams.thinking_start_tag = jinjaResult.thinking_start_tag
        if (typeof jinjaResult.thinking_end_tag === 'string')
          nativeParams.thinking_end_tag = jinjaResult.thinking_end_tag
        if (jinjaResult.chat_parser)
          nativeParams.chat_parser = jinjaResult.chat_parser
      } else if (formattedResult.type === 'llama-chat') {
        const llamaChatResult = formattedResult as FormattedChatResult
        nativeParams.prompt = llamaChatResult.prompt || ''
        if (llamaChatResult.has_media) {
          nativeParams.media_paths = llamaChatResult.media_paths
        }
      }
    } else {
      nativeParams.prompt = params.prompt || ''
    }

    if (!nativeParams.media_paths && params.media_paths) {
      nativeParams.media_paths = params.media_paths
    }

    if (params.response_format && !nativeParams.grammar) {
      const jsonSchema = getJsonSchema(params.response_format)
      if (jsonSchema) nativeParams.json_schema = JSON.stringify(jsonSchema)
    }

    // An empty prompt is valid for embedding-injection flows — e.g. TTS models
    // like Chatterbox T3 whose text is tokenized and embedded natively and
    // injected via the completion loop (getFormattedAudioCompletion returns
    // embedding:true for these). Only require a text prompt for a normal
    // token completion with neither embedding mode nor media inputs.
    if (!nativeParams.prompt && !params.embedding && !nativeParams.media_paths)
      throw new Error('Prompt is required')

    const { llamaCompletion } = getJsi()
    return llamaCompletion(this.id, nativeParams, callback)
  }

  stopCompletion(): Promise<void> {
    const { llamaStopCompletion } = getJsi()
    return llamaStopCompletion(this.id)
  }

  tokenize(
    text: string,
    {
      media_paths: mediaPaths,
    }: {
      media_paths?: string[]
    } = {},
  ): Promise<NativeTokenizeResult> {
    const { llamaTokenize } = getJsi()
    return llamaTokenize(this.id, text, mediaPaths)
  }

  detokenize(tokens: number[]): Promise<string> {
    const { llamaDetokenize } = getJsi()
    return llamaDetokenize(this.id, tokens)
  }

  embedding(
    text: string,
    params?: EmbeddingParams,
  ): Promise<NativeEmbeddingResult> {
    const { llamaEmbedding } = getJsi()
    return llamaEmbedding(this.id, text, params || {})
  }

  async rerank(
    query: string,
    documents: string[],
    params?: RerankParams,
  ): Promise<RerankResult[]> {
    const { llamaRerank } = getJsi()
    const results = await llamaRerank(this.id, query, documents, params || {})

    return results
      .map((result) => ({
        ...result,
        document: documents[result.index],
      }))
      .sort((a, b) => b.score - a.score)
  }

  async bench(
    pp: number,
    tg: number,
    pl: number,
    nr: number,
  ): Promise<BenchResult> {
    const { llamaBench } = getJsi()
    const result = await llamaBench(this.id, pp, tg, pl, nr)
    const parsed = JSON.parse(result)
    return {
      nKvMax: parsed.n_kv_max,
      nBatch: parsed.n_batch,
      nUBatch: parsed.n_ubatch,
      flashAttn: parsed.flash_attn,
      isPpShared: parsed.is_pp_shared,
      nGpuLayers: parsed.n_gpu_layers,
      nThreads: parsed.n_threads,
      nThreadsBatch: parsed.n_threads_batch,
      pp: parsed.pp,
      tg: parsed.tg,
      pl: parsed.pl,
      nKv: parsed.n_kv,
      tPp: parsed.t_pp,
      speedPp: parsed.speed_pp,
      tTg: parsed.t_tg,
      speedTg: parsed.speed_tg,
      t: parsed.t,
      speed: parsed.speed,
    }
  }

  async applyLoraAdapters(
    loraList: Array<{ path: string; scaled?: number }>,
  ): Promise<void> {
    const { llamaApplyLoraAdapters } = getJsi()
    const loraAdapters = normalizeLoraAdapters({ loraList })
    return llamaApplyLoraAdapters(this.id, loraAdapters)
  }

  async removeLoraAdapters(): Promise<void> {
    const { llamaRemoveLoraAdapters } = getJsi()
    return llamaRemoveLoraAdapters(this.id)
  }

  async getLoadedLoraAdapters(): Promise<
    Array<{ path: string; scaled?: number }>
  > {
    const { llamaGetLoadedLoraAdapters } = getJsi()
    return llamaGetLoadedLoraAdapters(this.id)
  }

  /**
   * Initialize multimodal support (vision/audio) with a projector model.
   * @param path - Path to the multimodal projector model file (mmproj)
   * @param use_gpu - Whether to use GPU for multimodal processing (default: true)
   * @param image_min_tokens - Minimum number of tokens for image input (for dynamic resolution models)
   * @param image_max_tokens - Maximum number of tokens for image input (for dynamic resolution models).
   *                           Lower values reduce memory usage and improve speed for high-resolution images.
   *                           Recommended: 256-512 for faster inference, up to 4096 for maximum detail.
   */
  async initMultimodal({
    path,
    use_gpu: useGpu,
    image_min_tokens: imageMinTokens,
    image_max_tokens: imageMaxTokens,
  }: {
    path: string
    use_gpu?: boolean
    image_min_tokens?: number
    image_max_tokens?: number
  }): Promise<boolean> {
    const { llamaInitMultimodal } = getJsi()
    if (path.startsWith('file://')) path = path.slice(7)
    return llamaInitMultimodal(this.id, {
      path,
      use_gpu: useGpu ?? true,
      image_min_tokens: imageMinTokens,
      image_max_tokens: imageMaxTokens,
    })
  }

  async isMultimodalEnabled(): Promise<boolean> {
    const { llamaIsMultimodalEnabled } = getJsi()
    return await llamaIsMultimodalEnabled(this.id)
  }

  async getMultimodalSupport(): Promise<{
    vision: boolean
    audio: boolean
  }> {
    const { llamaGetMultimodalSupport } = getJsi()
    return await llamaGetMultimodalSupport(this.id)
  }

  async releaseMultimodal(): Promise<void> {
    const { llamaReleaseMultimodal } = getJsi()
    return await llamaReleaseMultimodal(this.id)
  }

  /**
   * Attach a codec / vocoder GGUF to this context, enabling the TTS API.
   *
   * **Experimental:** the TTS API may change without a major version bump, and
   * output quality varies by model family and backend. See the "Tested models"
   * table in the README.
   */
  async initVocoder({
    path,
    n_batch: nBatch,
    use_gpu: useGpu,
  }: {
    path: string
    n_batch?: number
    // Offload the codec / codec_lm graphs (Mimi / S3G / depth decoder /
    // RVQ heads etc.) to GPU via codec.cpp's `ggml_backend_init_best`.
    // Defaults to true when the main context was created with
    // `n_gpu_layers > 0`, false otherwise; pass explicitly to override.
    use_gpu?: boolean
  }): Promise<boolean> {
    const { llamaInitVocoder } = getJsi()
    if (path.startsWith('file://')) path = path.slice(7)
    return await llamaInitVocoder(this.id, {
      path,
      n_batch: nBatch,
      use_gpu: useGpu,
    })
  }

  async isVocoderEnabled(): Promise<boolean> {
    const { llamaIsVocoderEnabled } = getJsi()
    return await llamaIsVocoderEnabled(this.id)
  }

  async getTTSCapabilities(): Promise<import('./tts').TTSCapabilities> {
    const { llamaGetTTSCapabilities } = getJsi()
    return await llamaGetTTSCapabilities(this.id)
  }

  /**
   * Build a formatted prompt for the loaded TTS model.
   *
   * Breaking change: takes an options object — the previous `(speaker, text)`
   * positional signature has been removed.
   *
   * - `prompt` — text to speak. Phonemized if `phonemizer` is supplied.
   * - `speaker` — built-in voice name (string), a structured speaker object
   *   (shape depends on the model family — see `OuteTTSSpeaker` /
   *   `NeuTTSSpeaker`), or `undefined` to fall back to the family default.
   * - `phonemizer` — optional `(text, language) => string | Promise<string>`.
   *   When set, `prompt` and `speaker.ref_text` (if missing `ref_phones`) go
   *   through it. Models that need phonemes (NeuTTS) get off-distribution
   *   text otherwise — caller's call.
   * - `language` — phonemizer hook hint; defaults to capabilities.defaultLanguage.
   */
  async getFormattedAudioCompletion(options: {
    prompt: string
    speaker?: string | LlamaSpeaker | SpeakerPayload
    phonemizer?: (
      text: string,
      language: string,
    ) => string | Promise<string>
    language?: string
  }): Promise<{
    prompt: string
    grammar?: string
    embedding: boolean
    flow: 'tokens' | 'continuous_embd' | ''
  }> {
    const { lookupVoice } = require('./tts-voices')
    const cap = await this.getTTSCapabilities()
    const language = options.language ?? (cap.defaultLanguage || 'en-us')

    // 1. phonemize input text (only if hook supplied — caller controls)
    let inputText = options.prompt
    if (cap.requiresPhonemes && options.phonemizer) {
      inputText = await Promise.resolve(
        options.phonemizer(options.prompt, language),
      )
    }

    const { llamaGetFormattedAudioCompletion } = getJsi()

    // 2. LlamaSpeaker handle path: pass empty speakerStr + the speaker id so
    //    native arms pending_speaker_id from the registry. Downstream
    //    completion / generateAudioCodes calls inject the speaker automatically.
    if (options.speaker instanceof LlamaSpeaker) {
      return llamaGetFormattedAudioCompletion(
        this.id,
        '',
        inputText,
        options.speaker.id,
      )
    }

    // 3. Otherwise resolve to a pre-baked speaker payload (an OuteTTSSpeaker /
    //    NeuTTSSpeaker config) and forward it to native as speaker JSON. A
    //    structured object is used directly; a string name — or `undefined`,
    //    which means 'default' — resolves against the built-in voice table,
    //    whose entries are themselves pre-baked payloads. Native never sees a
    //    voice-name key; voice resolution is JS-only.
    let payload: Record<string, any> | null
    if (options.speaker && typeof options.speaker === 'object') {
      payload = { ...options.speaker }
    } else {
      const name =
        typeof options.speaker === 'string' ? options.speaker : 'default'
      // The built-in OuteTTS voice is a legacy word/codes payload that only
      // fits the legacy / V0.3 prompt builders. OuteTTS 1.0 expects per-word
      // c1/c2 codes — resolving the legacy payload for it poisons the prompt
      // (codeless word blocks), so V1.0 falls back to speaker-less generation.
      const voice =
        cap.promptKind === 'outetts_v1_0'
          ? null
          : lookupVoice(cap.family, name, language)
      if (!voice && typeof options.speaker === 'string') {
        // Not a type check — this reports an unknown voice *value*, so Error
        // (not TypeError) is correct despite the enclosing typeof guard.
        // eslint-disable-next-line unicorn/prefer-type-error
        throw new Error(
          `Unknown built-in voice '${name}' for ${cap.family || 'this model'} (${language})`,
        )
      }
      payload = voice ? { ...voice } : null
    }

    // NeuTTS: phonemize the payload's ref_text → ref_phones if the caller
    // didn't already pre-phonemize it.
    if (
      payload &&
      cap.requiresPhonemes &&
      !payload.ref_phones &&
      payload.ref_text &&
      options.phonemizer
    ) {
      payload.ref_phones = await Promise.resolve(
        options.phonemizer(payload.ref_text, language),
      )
    }

    const speakerStr = payload ? JSON.stringify(payload) : ''
    return llamaGetFormattedAudioCompletion(this.id, speakerStr, inputText)
  }

  async decodeAudioTokens(tokens: number[]): Promise<Array<number>> {
    const { llamaDecodeAudioTokens } = getJsi()
    return await llamaDecodeAudioTokens(this.id, tokens)
  }

  /**
   * DEPRECATED: source-compat wrapper for codec_lm-AR TTS.
   *
   * As of the "one completion API" refactor, codec_lm-AR models (CSM /
   * Qwen3-TTS / MOSS-TTSD / MOSS-TTS-Realtime / Chatterbox) run through
   * the standard `completion` loop with `flow = 'tokens'` and
   * `embedding = true`.  The per-step codec_lm state machine that used
   * to live inside this call is now a hook on the completion loop
   * (`tryCodecLmAudioStep`); the codes get appended to
   * `result.audio_tokens` the same way OuteTTS / Soprano / NeuTTS do.
   *
   * This method still works — internally it just primes params +
   * speaker prefix, runs `completion`, and drains `audio_tokens` — but
   * new callers should skip it and use `completion()` +
   * `decodeAudioTokens` directly.
   *
   * `onFrame` (optional) fires after each AR step with that frame's
   * codes for streaming UIs. It is fire-and-forget — its return value
   * isn't read.
   */
  async generateAudioCodes(options: {
    prompt: string
    maxFrames?: number
    temperature?: number
    topP?: number
    topK?: number
    seed?: number
    onFrame?: (step: number, codes: number[]) => void
  }): Promise<{
    codes: number[]
    nCodebook: number
    nFrames: number
    stoppedOnEos: boolean
    aborted: boolean
  }> {
    const { llamaGenerateAudioCodes } = getJsi()
    const { onFrame, ...rest } = options
    const optsJson = JSON.stringify(rest)
    return await llamaGenerateAudioCodes(this.id, optsJson, onFrame)
  }

  async createSpeaker(config: {
    refAudio: Float32Array | number[]
    refAudioSampleRate: number
    refText?: string
    emotion?: number
    bake?: boolean
  }): Promise<LlamaSpeaker> {
    const { llamaCreateSpeaker } = getJsi()
    const pcm =
      config.refAudio instanceof Float32Array
        ? Array.from(config.refAudio)
        : config.refAudio
    const optsJson = JSON.stringify({
      pcm,
      inputSampleRate: config.refAudioSampleRate,
      refText: config.refText ?? '',
      bake: config.bake ?? false,
      ...(config.emotion !== undefined ? { emotion: config.emotion } : {}),
    })
    const h = await llamaCreateSpeaker(this.id, optsJson)
    return new LlamaSpeaker(this.id, h)
  }

  async decodeAudioEmbeddings(
    embeddings: number[],
    embeddingDim: number,
  ): Promise<Array<number>> {
    const { llamaDecodeAudioEmbeddings } = getJsi()
    return await llamaDecodeAudioEmbeddings(this.id, embeddings, embeddingDim)
  }

  async getAudioSampleRate(): Promise<number> {
    const { llamaGetAudioSampleRate } = getJsi()
    return await llamaGetAudioSampleRate(this.id)
  }

  async releaseVocoder(): Promise<void> {
    const { llamaReleaseVocoder } = getJsi()
    return await llamaReleaseVocoder(this.id)
  }

  /**
   * Clear the KV cache and reset conversation state
   * @param clearData If true, clears both metadata and tensor data buffers (slower). If false, only clears metadata (faster).
   * @returns Promise that resolves when cache is cleared
   *
   * Call this method between different conversations to prevent cache contamination.
   * Without clearing, the model may use cached context from previous conversations,
   * leading to incorrect or unexpected responses.
   *
   * For hybrid architecture models (e.g., LFM2), this is essential as they
   * use recurrent state that cannot be partially removed - only fully cleared.
   */
  async clearCache(clearData: boolean = false): Promise<void> {
    const { llamaClearCache } = getJsi()
    return llamaClearCache(this.id, clearData)
  }

  async release(): Promise<void> {
    const { llamaReleaseContext } = getJsi()
    return llamaReleaseContext(this.id)
  }
}

export async function toggleNativeLog(enabled: boolean): Promise<void> {
  await installJsi()
  const { llamaToggleNativeLog } = getJsi()
  return llamaToggleNativeLog(enabled, emitNativeLog)
}

export function addNativeLogListener(
  listener: (level: string, text: string) => void,
): { remove: () => void } {
  logListeners.push(listener)
  return {
    remove: () => {
      logListeners.splice(logListeners.indexOf(listener), 1)
    },
  }
}

export async function setContextLimit(limit: number): Promise<void> {
  await installJsi()
  const { llamaSetContextLimit } = getJsi()
  return llamaSetContextLimit(limit)
}

let contextIdCounter = 0
const contextIdRandom = () =>
  /* @ts-ignore */
  process.env.NODE_ENV === 'test' ? 0 : Math.floor(Math.random() * 100000)

const modelInfoSkip = [
  'tokenizer.ggml.tokens',
  'tokenizer.ggml.token_type',
  'tokenizer.ggml.merges',
  'tokenizer.ggml.scores',
]
export async function loadLlamaModelInfo(model: string): Promise<Object> {
  await installJsi()
  const { llamaModelInfo } = getJsi()
  let path = model
  if (path.startsWith('file://')) path = path.slice(7)
  return llamaModelInfo(path, modelInfoSkip)
}

const poolTypeMap = {
  none: 0,
  mean: 1,
  cls: 2,
  last: 3,
  rank: 4,
}

const normalizeFileUri = (path?: string) =>
  path?.startsWith('file://') ? path.slice(7) : path

const normalizeSpeculativeDraftPaths = (
  speculative?: NativeSpeculativeConfig,
): NativeSpeculativeConfig | undefined => {
  if (
    !speculative ||
    typeof speculative !== 'object' ||
    Array.isArray(speculative)
  ) {
    return speculative
  }

  const { draft } = speculative
  if (!draft || typeof draft !== 'object') {
    return speculative
  }

  return {
    ...speculative,
    draft: {
      ...draft,
      model: normalizeFileUri(draft.model),
      path: normalizeFileUri(draft.path),
      model_draft: normalizeFileUri(draft.model_draft),
      draft_model: normalizeFileUri(draft.draft_model),
    },
  }
}

type LoraAdapterInput = {
  path: string
  scaled?: number
}

const stripFilePrefix = (path: string) => path.replace(/^file:\/\//, '')

function normalizeLoraAdapters({
  lora,
  loraScaled,
  loraList,
}: {
  lora?: string
  loraScaled?: number
  loraList?: Array<LoraAdapterInput>
}): Array<LoraAdapterInput> {
  const adapters = new Map<string, LoraAdapterInput>()

  if (lora) {
    const normalizedPath = stripFilePrefix(lora)
    if (normalizedPath) {
      adapters.set(normalizedPath, {
        path: normalizedPath,
        scaled: loraScaled,
      })
    }
  }

  if (loraList) {
    loraList.forEach((adapter) => {
      const normalizedPath = stripFilePrefix(adapter.path)
      if (!normalizedPath) return
      adapters.set(normalizedPath, {
        path: normalizedPath,
        scaled: adapter.scaled,
      })
    })
  }

  return Array.from(adapters.values())
}

export async function getBackendDevicesInfo(): Promise<
  Array<NativeBackendDeviceInfo>
> {
  await installJsi()
  const { llamaGetBackendDevicesInfo } = getJsi()
  try {
    const jsonString = await llamaGetBackendDevicesInfo()
    return JSON.parse(jsonString as string)
  } catch (e) {
    console.warn(
      '[RNLlama] Failed to parse backend devices info, falling back to empty list',
      e,
    )
    return []
  }
}

export async function initLlama(
  {
    model,
    model_draft: modelDraft,
    draft_model: draftModelAlias,
    is_model_asset: isModelAsset,
    pooling_type: poolingType,
    lora,
    lora_scaled: loraScaled,
    lora_list: loraList,
    devices,
    ...rest
  }: ContextParams,
  onProgress?: (progress: number) => void,
): Promise<LlamaContext> {
  await installJsi()
  const { llamaInitContext } = getJsi()
  const path = normalizeFileUri(model) || model
  const draftPath = normalizeFileUri(modelDraft || draftModelAlias)
  const nativeRest = {
    ...rest,
    speculative: normalizeSpeculativeDraftPaths(rest.speculative),
  }

  const loraAdapters = normalizeLoraAdapters({
    lora,
    loraScaled,
    loraList,
  })

  const contextId = contextIdCounter + contextIdRandom()
  contextIdCounter += 1

  let lastProgress = 0
  const progressCallback = onProgress
    ? (progress: number) => {
        lastProgress = progress
        try {
          onProgress(progress)
        } catch (err) {
          console.warn('[RNLlama] onProgress callback failed', err)
        }
      }
    : undefined

  if (progressCallback) progressCallback(0)

  const poolType = poolTypeMap[poolingType as keyof typeof poolTypeMap]

  if (
    nativeRest.cache_type_k &&
    !validCacheTypes.includes(nativeRest.cache_type_k)
  ) {
    console.warn(
      `[RNLlama] initLlama: Invalid cache K type: ${nativeRest.cache_type_k}, falling back to f16`,
    )
    delete nativeRest.cache_type_k
  }
  if (
    nativeRest.cache_type_v &&
    !validCacheTypes.includes(nativeRest.cache_type_v)
  ) {
    console.warn(
      `[RNLlama] initLlama: Invalid cache V type: ${nativeRest.cache_type_v}, falling back to f16`,
    )
    delete nativeRest.cache_type_v
  }

  let filteredDevs: Array<string> = []
  if (Array.isArray(devices)) {
    filteredDevs = [...devices]
    const backendDevices = await getBackendDevicesInfo()

    if (Platform.OS === 'android' && devices.includes('HTP*')) {
      const htpDevices = backendDevices
        .filter((d) => d.deviceName.startsWith('HTP'))
        .map((d) => d.deviceName)
      filteredDevs = filteredDevs.reduce((acc, dev) => {
        if (dev.startsWith('HTP*')) {
          acc.push(...htpDevices)
        } else if (!dev.startsWith('HTP')) {
          acc.push(dev)
        }
        return acc
      }, [] as Array<string>)
    }
  }

  const {
    gpu,
    devices: usedDevices,
    reasonNoGPU,
    model: modelDetails,
    androidLib,
    systemInfo,
  } = await llamaInitContext(
    contextId,
    {
      model: path,
      is_model_asset: !!isModelAsset,
      use_progress_callback: !!progressCallback,
      pooling_type: poolType,
      ...(loraAdapters.length > 0 ? { lora_list: loraAdapters } : {}),
      ...(draftPath ? { model_draft: draftPath } : {}),
      devices: filteredDevs.length > 0 ? filteredDevs : undefined,
      ...nativeRest,
    },
    progressCallback,
  )

  if (progressCallback && lastProgress < 100) progressCallback(100)

  return new LlamaContext({
    contextId,
    gpu,
    devices: usedDevices,
    reasonNoGPU,
    model: modelDetails,
    androidLib,
    systemInfo,
  })
}

export async function releaseAllLlama(): Promise<void> {
  if (!isJsiInstalled) return
  const { llamaReleaseAllContexts } = getJsi()
  return llamaReleaseAllContexts()
}

export const BuildInfo = {
  number: BUILD_NUMBER,
  commit: BUILD_COMMIT,
}