import { describe, expect, it } from 'vitest'
import {
  assignSamples,
  buildCohortTree,
  filterSamples,
  type ManualSample,
} from '../../components/SmartGroupManual.cohorts'

/** Mirrors the design prototype's seeded default sample list. */
function seededSamples(): ManualSample[] {
  const blank = { group: null, sub: null, sub2: null, color: null } as const
  const seed: Record<string, Partial<ManualSample>> = {
    Pt001_TumorA_d7_rep1: { group: 'Responder', sub: 'Day 7', sub2: 'Batch 1', color: 'c-teal' },
    Pt001_TumorA_d14_rep1: { group: 'Responder', sub: 'Day 14', color: 'c-teal' },
    Pt002_TumorB_d7_rep1: { group: 'Non-responder', sub: 'Day 7', sub2: 'Batch 2', color: 'c-rose' },
  }
  return ['Pt001_TumorA_d7_rep1', 'Pt001_TumorA_d14_rep1', 'Pt002_TumorB_d7_rep1', 'QC_pool_01'].map(
    name => ({ name, ...blank, ...(seed[name] ?? {}) }),
  )
}

describe('buildCohortTree', () => {
  const tree = buildCohortTree(seededSamples())

  it('should ignore unassigned samples', () => {
    expect(tree.size).toBe(2)
    expect(tree.has('Responder')).toBe(true)
    expect(tree.has('Non-responder')).toBe(true)
  })

  it('should total Responder across its two subgroups', () => {
    const responder = tree.get('Responder')!
    expect(responder.total).toBe(2)
    expect(responder.subs.get('Day 7')!.total).toBe(1)
    expect(responder.subs.get('Day 7')!.subs2.get('Batch 1')).toBe(1)
    expect(responder.subs.get('Day 14')!.total).toBe(1)
    expect(responder.subs.get('Day 14')!.subs2.get('—')).toBe(1)
  })

  it('should nest Non-responder under Day 7 / Batch 2', () => {
    const nonResponder = tree.get('Non-responder')!
    expect(nonResponder.total).toBe(1)
    expect(nonResponder.subs.get('Day 7')!.subs2.get('Batch 2')).toBe(1)
  })
})

describe('filterSamples', () => {
  const samples = seededSamples()

  it('should match names case-insensitively on a substring query', () => {
    const matches = filterSamples(samples, 'tumorb', 'all')
    expect(matches.map(m => m.sample.name)).toEqual(['Pt002_TumorB_d7_rep1'])
  })

  it('should preserve the original index for each match', () => {
    const matches = filterSamples(samples, 'QC', 'all')
    expect(matches).toEqual([{ sample: samples[3], index: 3 }])
  })

  it('should show only unassigned samples when filtered to unassigned', () => {
    const matches = filterSamples(samples, '', 'unassigned')
    expect(matches.map(m => m.sample.name)).toEqual(['QC_pool_01'])
  })

  it('should show only assigned samples when filtered to assigned', () => {
    const matches = filterSamples(samples, '', 'assigned')
    expect(matches).toHaveLength(3)
    expect(matches.every(m => !!m.sample.group)).toBe(true)
  })
})

describe('assignSamples', () => {
  const samples = seededSamples()

  it('should assign group and colour to the selected indices only', () => {
    const next = assignSamples(samples, new Set([3]), {
      grp: 'Pooled',
      sub: '',
      sub2: '',
      color: 'c-slate',
    })
    expect(next[3]).toMatchObject({ group: 'Pooled', color: 'c-slate' })
    expect(next[0]).toEqual(samples[0])
  })

  it('should null out blank subgroup and sub-subgroup values', () => {
    const next = assignSamples(samples, new Set([3]), {
      grp: '  Pooled  ',
      sub: '   ',
      sub2: '',
      color: 'c-slate',
    })
    expect(next[3].group).toBe('Pooled')
    expect(next[3].sub).toBeNull()
    expect(next[3].sub2).toBeNull()
  })
})