import type { SampleClass, SampleClassKind, ClassDisposition } from '../../types/auto-group'
import { classKey } from './classKey'
import vocabData from './vocab.json'

export interface ClassVocab {
  blankPattern: RegExp
  iqcPattern: RegExp
  eqcPattern: RegExp
  standardPattern: RegExp
  standardAllow: Set<string>
  spikeInPattern: RegExp
  stockPattern: RegExp
  /**
   * Matrix lookup: lowercase token → canonical display name. Matching is
   * case-insensitive, and several aliases can collapse to one display name
   * (cell → Cells, medium → Medium).
   */
  matrixVocab: Map<string, string>
  /** Lowercase polarity tokens → POS / NEG. */
  ionizationModes: Map<string, 'POS' | 'NEG'>
  /** Lowercase tokens that flag the next token as the organ subKind. */
  tissueParentTokens: Set<string>
}

type VocabFile = typeof vocabData
type IonMode = 'POS' | 'NEG'

function compileVocab(data: VocabFile): ClassVocab {
  return {
    blankPattern: new RegExp(data.patterns.blank, 'i'),
    iqcPattern: new RegExp(data.patterns.iqc, 'i'),
    eqcPattern: new RegExp(data.patterns.eqc, 'i'),
    standardPattern: new RegExp(data.patterns.standard, 'i'),
    spikeInPattern: new RegExp(data.patterns.spikeIn, 'i'),
    stockPattern: new RegExp(data.patterns.stock, 'i'),
    standardAllow: new Set(data.standardAllow),
    matrixVocab: new Map(Object.entries(data.matrixVocab)),
    ionizationModes: new Map(
      Object.entries(data.ionizationModes).map(([k, v]) => [k, v as IonMode]),
    ),
    tissueParentTokens: new Set(data.tissueParentTokens),
  }
}

export const classVocab: ClassVocab = compileVocab(vocabData)

function matrixCanonical(token: string, vocab: ClassVocab): string | undefined {
  return vocab.matrixVocab.get(token.toLowerCase())
}

function isTissueParentToken(token: string, vocab: ClassVocab): boolean {
  return vocab.tissueParentTokens.has(token.toLowerCase())
}

export interface DetectClassOptions {
  vocab?: Partial<ClassVocab>
  /** Per-sample explicit `sample_type` hint (e.g. from experiment design data). */
  sampleTypeHints?: (string | undefined)[]
  /** Number of leading tokens that all samples share (`TokenizedBatch.sharedPrefix`). */
  sharedPrefix?: number
}

/** A token that names a QC-type class (blank, QC, standard, spike-in, stock). */
export function isClassMarker(token: string, vocab: ClassVocab = classVocab): boolean {
  return (
    vocab.blankPattern.test(token) || vocab.iqcPattern.test(token) || vocab.eqcPattern.test(token)
    || vocab.standardPattern.test(token) || vocab.standardAllow.has(token)
    || vocab.spikeInPattern.test(token) || vocab.stockPattern.test(token)
  )
}

/** Kinds that form experimental groups; all other kinds are QC-type overlay classes. */
export function isExperimentalKind(kind: SampleClassKind): boolean {
  return kind === 'biological' || kind === 'unknown'
}

const DEFAULT_DISPOSITION: Record<SampleClassKind, ClassDisposition> = {
  biological: 'group',
  iqc: 'overlay',
  eqc: 'overlay',
  standard: 'overlay',
  'spike-in': 'overlay',
  stock: 'overlay',
  blank: 'overlay',
  unknown: 'group',
}

function classifyTokens(
  tokens: string[],
  vocab: ClassVocab,
  qcFrom: number,
): { kind: SampleClassKind; subKind?: string; tagPositions: number[] } {
  // First-match-wins ordering: Blank → IQC → EQC → Standard → SpikeIn → Stock → Biological → Unknown.
  // QC kinds are read from `qcFrom`: a token that all names share (the
  // experiment title) cannot mark one sample as QC.
  for (let i = qcFrom; i < tokens.length; i++) {
    if (vocab.blankPattern.test(tokens[i])) return { kind: 'blank', tagPositions: [i] }
  }
  for (let i = qcFrom; i < tokens.length; i++) {
    if (vocab.iqcPattern.test(tokens[i])) return { kind: 'iqc', tagPositions: [i] }
  }
  for (let i = qcFrom; i < tokens.length; i++) {
    if (vocab.eqcPattern.test(tokens[i])) return { kind: 'eqc', tagPositions: [i] }
  }
  for (let i = qcFrom; i < tokens.length; i++) {
    if (vocab.standardPattern.test(tokens[i]) || vocab.standardAllow.has(tokens[i])) {
      return { kind: 'standard', tagPositions: [i] }
    }
  }
  for (let i = qcFrom; i < tokens.length; i++) {
    if (vocab.spikeInPattern.test(tokens[i])) return { kind: 'spike-in', tagPositions: [i] }
  }
  for (let i = qcFrom; i < tokens.length; i++) {
    if (vocab.stockPattern.test(tokens[i])) return { kind: 'stock', tagPositions: [i] }
  }
  for (let i = 0; i < tokens.length; i++) {
    const canonical = matrixCanonical(tokens[i], vocab)
    if (canonical) {
      const isTissueParent = isTissueParentToken(tokens[i], vocab)
      const tagPositions = [i]
      if (isTissueParent) {
        // "Tissues"/"Tissue" is the parent group; tag the next token (organ name) as well.
        if (i + 1 < tokens.length) tagPositions.push(i + 1)
        return { kind: 'biological', subKind: 'Tissues', tagPositions }
      }
      // Non-parent matrix token: tag any immediately preceding "Tissues"/"Tissue" parent.
      if (i > 0 && isTissueParentToken(tokens[i - 1], vocab)) {
        tagPositions.unshift(i - 1)
        return { kind: 'biological', subKind: 'Tissues', tagPositions }
      }
      return { kind: 'biological', subKind: canonical, tagPositions }
    }
  }
  return { kind: 'unknown', tagPositions: [] }
}

/**
 * LC-MS batches run positive- and negative-ion-mode acquisitions separately,
 * and the two produce different feature tables that should never be grouped
 * together. When a POS / NEG / positive / negative token is present we suffix
 * it to the class label so the schema builder treats each mode as its own class.
 */
function ionizationModeFor(
  tokens: string[],
  vocab: ClassVocab,
): { mode: IonMode; position: number } | undefined {
  for (let i = 0; i < tokens.length; i++) {
    const mode = vocab.ionizationModes.get(tokens[i].toLowerCase())
    if (mode) return { mode, position: i }
  }
  return undefined
}

function labelFor(kind: SampleClassKind, subKind?: string): string {
  const base: Record<SampleClassKind, string> = {
    biological: 'Biological',
    iqc: 'IQC',
    eqc: 'EQC',
    standard: 'Standard',
    'spike-in': 'SpikeIn',
    stock: 'Stock',
    blank: 'Blank',
    unknown: 'Unknown',
  }
  return subKind ? `${base[kind]} / ${subKind}` : base[kind]
}

/** Classify each tokenized sample into a SampleClass. */
export function detectClass(
  tokenizedSamples: string[][],
  options: DetectClassOptions,
): SampleClass[] {
  const vocab = { ...classVocab, ...(options.vocab ?? {}) }
  const hints = options.sampleTypeHints ?? []
  const buckets = new Map<string, SampleClass>()

  for (let i = 0; i < tokenizedSamples.length; i++) {
    const hint = hints[i]?.toLowerCase()
    let kind: SampleClassKind
    let subKind: string | undefined
    let tagPositions: number[]

    if (hint === 'qc') {
      kind = 'iqc'
      tagPositions = []
    } else if (hint === 'blank') {
      kind = 'blank'
      tagPositions = []
    } else {
      const r = classifyTokens(tokenizedSamples[i], vocab, options.sharedPrefix ?? 0)
      kind = r.kind
      subKind = r.subKind
      tagPositions = r.tagPositions
    }

    // Split classes by ionization mode (POS / NEG) so the two acquisitions get
    // independent schemas even when they share a matrix and sample-type vocab.
    const mode = ionizationModeFor(tokenizedSamples[i], vocab)
    if (mode) {
      subKind = subKind ? `${subKind} / ${mode.mode}` : mode.mode
      if (!tagPositions.includes(mode.position)) tagPositions.push(mode.position)
    }

    const key = classKey({ kind, subKind })
    const bucket = buckets.get(key)
    if (bucket) {
      bucket.members.push(i)
      for (const p of tagPositions) {
        if (!bucket.classTagPositions.includes(p)) bucket.classTagPositions.push(p)
      }
    } else {
      buckets.set(key, {
        kind,
        subKind,
        label: labelFor(kind, subKind),
        members: [i],
        classTagPositions: [...tagPositions],
        disposition: DEFAULT_DISPOSITION[kind],
      })
    }
  }

  return [...buckets.values()]
}