import fs from 'node:fs'
import { Tokenizer } from '@huggingface/tokenizers'
import { chunkedPrefixTokens, fimPrompt, headTokens, repoHeader, wantsSingleLine } from '../../shared/prompt.ts'
import type { ContextMode, DaemonConfig, ModelProfile } from './config.ts'

export interface FimRequest {
  requestId?: string
  /** Workspace folder name; Stage C was trained with the repo directory name here. */
  repo: string
  /** Path relative to the workspace folder, '/'-separated. */
  path: string
  /** Absolute workspace folder on the daemon's machine; enables symbols mode and the this.-repair. */
  workspaceRoot?: string
  languageId?: string
  prefix: string
  suffix: string
  /** Absolute line number of prefix's first line, if the client sent a truncated document. */
  prefixLineOffset?: number
  mode?: ContextMode
  /** Override automatic single-line detection. */
  singleLine?: boolean
  maxLines?: number
}

export interface BuiltPrompt {
  ids: number[]
  mode: ContextMode
  singleLine: boolean
  prefixStartLine: number
  tokenizeMs: number
}

export const STOP_STRINGS = ['<|endoftext|>', '<|fim_prefix|>', '<|fim_suffix|>', '<|fim_middle|>', '<|fim_pad|>', '<|file_sep|>', '<|repo_name|>', '<|im_start|>', '<|im_end|>']

export class PromptBuilder {
  private encodeFn: (s: string) => number[]
  private fim: { prefix: number; suffix: number; middle: number }
  private header: { repoName: number; fileSep: number }
  private config: DaemonConfig

  constructor(config: DaemonConfig) {
    this.config = config
    const tokenizer = new Tokenizer(
      JSON.parse(fs.readFileSync(config.tokenizer.json, 'utf8')),
      JSON.parse(fs.readFileSync(config.tokenizer.config, 'utf8')),
    )
    this.encodeFn = s => tokenizer.encode(s, { add_special_tokens: false }).ids
    const id = (t: string) => {
      const v = tokenizer.token_to_id(t)
      if (v === undefined) throw new Error(`tokenizer is missing ${t}`)
      return v
    }
    this.fim = { prefix: id('<|fim_prefix|>'), suffix: id('<|fim_suffix|>'), middle: id('<|fim_middle|>') }
    this.header = { repoName: id('<|repo_name|>'), fileSep: id('<|file_sep|>') }
  }

  encode(s: string) { return this.encodeFn(s) }

  /**
   * Builds the prompt in the model's trained layout. `symbols` context is passed
   * in by the caller (empty until the symbol provider exists, which degrades to `simple`).
   */
  build(req: FimRequest, profile: ModelProfile, mode: ContextMode, symbols?: string): BuiltPrompt {
    const t0 = performance.now()
    const p = this.config.prompt
    const { tokens: prefix, startLine } = chunkedPrefixTokens(req.prefix, p.prefixTokens, this.encodeFn, p.prefixChunkLines, req.prefixLineOffset ?? 0)
    const suffix = headTokens(req.suffix, p.suffixTokens, this.encodeFn)
    const header = mode === 'raw' ? [] : repoHeader(req.repo, req.path, mode === 'symbols' ? symbols : undefined, this.encodeFn, this.header)
    return {
      ids: [...header, ...fimPrompt(profile.order, prefix, suffix, this.fim)],
      mode,
      singleLine: req.singleLine ?? wantsSingleLine(req.suffix),
      prefixStartLine: startLine,
      tokenizeMs: performance.now() - t0,
    }
  }
}