import { describe, it, expect } from 'vitest'
import { composeGroups } from '../../../composables/autoGroup/compose'
import type { ClassSchema, SampleClass } from '../../../types/auto-group'
import type { TreeNode } from '../../../types'

const tissues: ClassSchema = {
  classKind: 'biological',
  subKind: 'Tissues',
  columns: [
    { index: 0, name: 'Organ', sourceIndices: [0], uniqueValues: ['Kidney', 'Liver'], cardinality: 2, role: 'factor' },
    { index: 1, name: 'Subject', sourceIndices: [1], uniqueValues: ['M1', 'M2'], cardinality: 2, role: 'factor' },
    { index: 2, name: 'Injection #', sourceIndices: [2], uniqueValues: ['001', '002', '003'], cardinality: 3, role: 'run-order' },
  ],
  groupBy: [0, 1],
}
const tissuesClass: SampleClass = {
  kind: 'biological',
  subKind: 'Tissues',
  label: 'Biological / Tissues',
  members: [0, 1, 2],
  classTagPositions: [],
  disposition: 'group',
}
const input = {
  tokenizedSamples: [
    ['Kidney', 'M1', '001'],
    ['Kidney', 'M2', '002'],
    ['Liver', 'M1', '003'],
  ],
  sampleNames: ['k1', 'k2', 'l1'],
  schemas: { 'biological:Tissues': tissues },
  classes: [tissuesClass],
}

function child(node: TreeNode | undefined, label: string): TreeNode | undefined {
  return node?.children?.find(c => c.label === label)
}

describe('composeGroups groupTree (nested hierarchy)', () => {
  it('nests class → layer1 → layer2 → sample leaves with sample-count badges', () => {
    const tree = composeGroups(input).groupTree!
    expect(tree).toHaveLength(1)
    const root = tree[0]
    expect(root.label).toBe('Biological / Tissues')
    expect(root.badge).toBe(3) // total samples under the class

    const kidney = child(root, 'Kidney')
    const liver = child(root, 'Liver')
    expect(kidney?.badge).toBe(2)
    expect(liver?.badge).toBe(1)

    const kidneyM1 = child(kidney, 'M1')
    const kidneyM2 = child(kidney, 'M2')
    expect(kidneyM1?.badge).toBe(1)
    expect(kidneyM2?.badge).toBe(1)

    // Leaf = the actual sample; injection number carried as leaf metadata + badge.
    const leaf = kidneyM1?.children?.[0]
    expect(leaf?.label).toBe('k1')
    expect(leaf?.children).toBeUndefined()
    expect(leaf?.metadata?.injection).toBe('001')
    expect(leaf?.badge).toBe('001')
  })

  it('gives every node a stable unique id derived from its path', () => {
    const tree = composeGroups(input).groupTree!
    const ids = new Set<string>()
    const walk = (n: TreeNode) => { ids.add(n.id); n.children?.forEach(walk) }
    tree.forEach(walk)
    // class + Kidney + Liver + M1(×2 paths) + M2 + 3 leaves = 9 distinct ids
    expect(ids.size).toBe(9)
  })

  it('puts samples directly under the class node when no columns are grouped', () => {
    const flatSchema: ClassSchema = { ...tissues, groupBy: [] }
    const tree = composeGroups({ ...input, schemas: { 'biological:Tissues': flatSchema } }).groupTree!
    const root = tree[0]
    expect(root.badge).toBe(3)
    expect(root.children?.map(c => c.label)).toEqual(['k1', 'k2', 'l1'])
  })

  it('excludes overlay/QC classes from the experimental groupTree', () => {
    const iqcSchema: ClassSchema = {
      classKind: 'iqc',
      columns: [{ index: 0, name: 'IQC', sourceIndices: [0], uniqueValues: ['IQC'], cardinality: 1, role: 'class-tag' }],
      groupBy: [],
    }
    const iqcClass: SampleClass = {
      kind: 'iqc', label: 'IQC', members: [3], classTagPositions: [0], disposition: 'overlay',
    }
    const tree = composeGroups({
      tokenizedSamples: [...input.tokenizedSamples, ['IQC']],
      sampleNames: [...input.sampleNames, 'qc1'],
      schemas: { 'biological:Tissues': tissues, iqc: iqcSchema },
      classes: [tissuesClass, iqcClass],
    }).groupTree!
    expect(tree.map(n => n.label)).toEqual(['Biological / Tissues'])
  })
})