import type { SampleClass } from '../../types/auto-group'
import { isExperimentalKind } from './classify'

/**
 * Token-length pre-grouping.
 *
 * After type classification (`detectClass`), members of one class can still have
 * different token counts — e.g. a Plasma batch where some names carry an extra
 * descriptor field. Splitting each class so every member shares the same token
 * length is the "group same token-length sample, then extract tokens inside the
 * group" step: it guarantees `buildClassSchema` sees a rectangular token matrix,
 * so column index N means the same field for every sample in the group.
 *
 * Each output class carries `tokenLength`; `classKey` folds that into the schema
 * key so length-split classes never collide. The label only gains a `· N fields`
 * suffix when a (kind, subKind) genuinely spans more than one length, keeping the
 * common (uniform) case visually identical to before.
 */
export function splitByTokenLength(
  classes: SampleClass[],
  tokenized: string[][],
): SampleClass[] {
  const out: SampleClass[] = []

  for (const cls of classes) {
    // QC-kind classes (QC, blank, standard) are one group named by the class
    // label and never group by column, so they need no rectangular matrix.
    if (!isExperimentalKind(cls.kind)) {
      out.push(cls)
      continue
    }

    // Bucket members by token count, preserving first-seen order of lengths.
    const byLength = new Map<number, number[]>()
    const lengthOrder: number[] = []
    for (const m of cls.members) {
      const len = tokenized[m]?.length ?? 0
      const bucket = byLength.get(len)
      if (bucket) {
        bucket.push(m)
      } else {
        byLength.set(len, [m])
        lengthOrder.push(len)
      }
    }

    const multipleLengths = lengthOrder.length > 1
    for (const len of lengthOrder) {
      out.push({
        ...cls,
        members: byLength.get(len)!,
        tokenLength: len,
        // Drop class-tag positions that fall outside this split's token range.
        classTagPositions: cls.classTagPositions.filter(p => p < len),
        label: multipleLengths ? `${cls.label} · ${len} fields` : cls.label,
      })
    }
  }

  return out
}