oc-ls-stats/src/stats.ts

export function backoffDelayMs(failures: number, baseMs: number = 1000, maxMs: number = 10000): number {
  return Math.min(maxMs, baseMs * 2 ** (failures - 1))
}

export interface SlotState {
  id: number
  is_processing: boolean
  n_prompt_tokens: number
  next_token: { n_decoded: number }[]
  id_task: number
}

export interface TrackerState {
  prevNdBySlot: Record<number, { prevNd: number; idTask: number; classification: "TG" | "PP" | "neither" }>
  prevPromptTokensBySlot: Record<number, { tokens: number; idTask: number }>
  isPrefilling: boolean
  isGenerating: boolean
  prefillSlotId: number | null
  prefillCapturedTokens: number | null
  prefillStartAt: number
  lastPrefillRate: number
  generateSlotId: number | null
  generatePrevNd: number
  generateStartAt: number
  lastGeneratedTps: number
}

export interface ClassificationResult {
  slotId: number
  nd: number
  prevNd: number
  nt: number
  prevNt: number
  classification: "TG" | "PP" | "neither"
}

export interface PollResult {
  classifications: ClassificationResult[]
  isPrefilling: boolean
  isGenerating: boolean
  lastPrefillRate: number
  lastGeneratedTps: number
  ppEnded: boolean
  tgEnded: boolean
}

export function classifySlot(
  slot: SlotState,
  tracker: TrackerState,
): ClassificationResult {
  const slotId = slot.id
  const nd = slot.next_token?.[0]?.n_decoded ?? 0
  const nt = slot.n_prompt_tokens ?? 0
  const idTask = slot.id_task ?? 0

  const ndEntry = tracker.prevNdBySlot[slotId]
  const ntEntry = tracker.prevPromptTokensBySlot[slotId]
  
  // Check if we have valid baseline for this slot/task
  const hasBaseline = ndEntry && ndEntry.idTask === idTask && ntEntry && ntEntry.idTask === idTask
  const baselineNd = hasBaseline ? ndEntry.prevNd : nd
  const baselineNt = hasBaseline ? ntEntry.tokens : nt
  
  // Check if values have reset (dropped below baseline) - indicates slot reuse or new request
  const hasReset = hasBaseline && (nd < baselineNd || nt < baselineNt)

  let classification: "TG" | "PP" | "neither" = "neither"
  let prevNd = nd
  let prevNt = nt

  if (!hasBaseline || hasReset) {
    // First time seeing this slot/task, or values have reset
    // Use heuristic: if n_decoded > 0, we likely caught it mid-generation
    if (nd > 0) {
      classification = "TG"
    } else {
      classification = "PP"
    }
  } else {
    // Compare against baseline values
    if (nd > baselineNd) {
      classification = "TG"
    } else if (nt > baselineNt) {
      classification = "PP"
    } else {
      // Neither has increased from baseline - persist previous classification
      // This handles frequent polling where we don't see increases every poll
      const prevClassification = ndEntry?.classification
      if (prevClassification) {
        classification = prevClassification
      }
    }
    prevNd = baselineNd
    prevNt = baselineNt
  }

  return { slotId, nd, prevNd, nt, prevNt, classification }
}

export function processPoll(
  slots: SlotState[],
  tracker: TrackerState,
  now: number = Date.now(),
): PollResult {
  const processingSlots = slots.filter(s => s.is_processing)
  const processingSlotIds = new Set(processingSlots.map(s => s.id))

  // Clean up tracking state for slots that went idle
  for (const key of Object.keys(tracker.prevNdBySlot)) {
    const slotId = Number(key)
    if (!processingSlotIds.has(slotId)) {
      delete tracker.prevNdBySlot[slotId]
    }
  }
  for (const key of Object.keys(tracker.prevPromptTokensBySlot)) {
    const slotId = Number(key)
    if (!processingSlotIds.has(slotId)) {
      delete tracker.prevPromptTokensBySlot[slotId]
    }
  }

  // Classify each processing slot
  const classifications: ClassificationResult[] = []
  let prefillSlot: ClassificationResult | null = null
  let generateSlot: ClassificationResult | null = null

  for (const slot of processingSlots) {
    const result = classifySlot(slot, tracker)
    classifications.push(result)

    if (result.classification === "TG") {
      generateSlot = result
    } else if (result.classification === "PP") {
      prefillSlot = result
    }

    // Update tracker state - store baseline values and classification
    tracker.prevNdBySlot[slot.id] = { prevNd: result.nd, idTask: slot.id_task, classification: result.classification }
    tracker.prevPromptTokensBySlot[slot.id] = { tokens: result.nt, idTask: slot.id_task }
  }

  let anyPrefilling = false
  let anyGenerating = false
  let ppEnded = false
  let tgEnded = false

  // Handle PP phase
  if (prefillSlot) {
    anyPrefilling = true
    const nt = prefillSlot.nt
    if (tracker.prefillSlotId === prefillSlot.slotId) {
      const dt = (now - tracker.prefillStartAt) / 1000
      const delta = nt - (tracker.prefillCapturedTokens ?? 0)
      if (dt > 0 && delta > 0) {
        tracker.lastPrefillRate = delta / dt
      }
    } else {
      tracker.prefillSlotId = prefillSlot.slotId
      tracker.prefillCapturedTokens = nt
      tracker.prefillStartAt = now
      tracker.lastPrefillRate = 0
    }
  } else if (processingSlotIds.size > 0 && !generateSlot) {
    anyPrefilling = tracker.isPrefilling
  } else if (tracker.isPrefilling) {
    tracker.prefillSlotId = null
    tracker.prefillCapturedTokens = null
    tracker.prefillStartAt = 0
    ppEnded = true
  }

  // Handle TG phase
  if (generateSlot) {
    anyGenerating = true
    const nd = generateSlot.nd
    if (tracker.generateSlotId === generateSlot.slotId) {
      const dt = (now - tracker.generateStartAt) / 1000
      const delta = nd - tracker.generatePrevNd
      if (dt > 0 && delta > 0) {
        tracker.lastGeneratedTps = delta / dt
      }
      tracker.generatePrevNd = nd
      tracker.generateStartAt = now
    } else {
      tracker.generateSlotId = generateSlot.slotId
      tracker.generatePrevNd = nd
      tracker.generateStartAt = now
      tracker.lastGeneratedTps = 0
    }
  } else if (processingSlotIds.size > 0 && !prefillSlot) {
    anyGenerating = tracker.isGenerating
  } else if (tracker.isGenerating) {
    tracker.generateSlotId = null
    tracker.generatePrevNd = 0
    tracker.generateStartAt = 0
    tracker.lastGeneratedTps = 0
    tgEnded = true
  }

  tracker.isPrefilling = anyPrefilling
  tracker.isGenerating = anyGenerating

  return {
    classifications,
    isPrefilling: tracker.isPrefilling,
    isGenerating: tracker.isGenerating,
    lastPrefillRate: tracker.lastPrefillRate,
    lastGeneratedTps: tracker.lastGeneratedTps,
    ppEnded,
    tgEnded,
  }
}

export function stripBaseUrlPath(baseURL: string): string {
  try {
    const url = new URL(baseURL)
    return `${url.origin}`
  } catch {
    return baseURL.replace(/\/v1(\/.*)?$/, "").replace(/\/$/, "") || baseURL
  }
}

export function extractProviderUrls(config: unknown): string[] {
  if (!config || typeof config !== "object") return []
  const obj = config as Record<string, unknown>
  const provider = obj.provider
  if (!provider || typeof provider !== "object") return []

  const urls: string[] = []
  for (const [key, val] of Object.entries(provider as Record<string, unknown>)) {
    if (val && typeof val === "object") {
      const opts = (val as Record<string, unknown>).options
      if (opts && typeof opts === "object") {
        const baseURL = (opts as Record<string, unknown>).baseURL || (opts as Record<string, unknown>).base_url
        if (typeof baseURL === "string" && key.toLowerCase().includes("llama") && !key.toLowerCase().includes("ollama")) {
          urls.push(stripBaseUrlPath(baseURL))
        }
      }
    }
  }
  return urls
}

export function extractProviderUrlsFromList(providers: unknown): string[] {
  if (!Array.isArray(providers)) return []
  const urls: string[] = []
  for (const entry of providers) {
    if (!entry || typeof entry !== "object") continue
    const provider = entry as Record<string, unknown>
    const id = typeof provider.id === "string" ? provider.id.toLowerCase() : ""
    const name = typeof provider.name === "string" ? provider.name.toLowerCase() : ""
    const hay = `${id} ${name}`
    if (!hay.includes("llama") || hay.includes("ollama")) continue
    const opts = provider.options
    if (!opts || typeof opts !== "object") continue
    const baseURL = (opts as Record<string, unknown>).baseURL || (opts as Record<string, unknown>).base_url
    if (typeof baseURL === "string") {
      urls.push(stripBaseUrlPath(baseURL))
    }
  }
  return urls
}

export function resolveServerUrls(options: unknown, providers: unknown, config: unknown): string[] {
  const opts = options && typeof options === "object" ? (options as Record<string, unknown>) : undefined
  const server = opts?.server
  const explicit = typeof server === "string" && server.trim() !== "" ? stripBaseUrlPath(server) : undefined
  if (explicit && /^https?:\/\//.test(explicit)) {
    return [explicit]
  }
  if (server !== undefined) {
    console.warn("[oc-ls-stats] ignoring invalid 'server' plugin option, falling back to provider detection")
  }
  const detected = extractProviderUrlsFromList(providers)
  if (detected.length > 0) return detected
  const legacy = extractProviderUrls(config)
  return legacy.length > 0 ? legacy : ["http://localhost:8080"]
}