import { describe, it, expect } from 'vitest';
import {
  applyDatasetKind,
  initialNnunetState,
  nnunetReducer,
  NNUNET_STEPS,
  type NnunetWizardState,
} from './wizardState';
import type { SourceFile } from './core';

const base = (): NnunetWizardState => structuredClone(initialNnunetState);
const file = (over: Partial<SourceFile> = {}): SourceFile => ({
  id: Math.random().toString(36).slice(2),
  filename: 'la_003_0000.nii.gz',
  sizeBytes: 1,
  role: 'unknown',
  split: 'train',
  ...over,
});

describe('applyDatasetKind', () => {
  it('volume3d from an RGB dataset resets to a single CT channel + a 3D ending', () => {
    const rgb = applyDatasetKind(base().meta, 'rgb2d');
    const vol = applyDatasetKind(rgb, 'volume3d');
    expect(vol.dim).toBe('3d');
    expect(vol.rgb).toBe(false);
    expect(vol.fileEnding).toBe('.nii.gz');
    expect(vol.channels).toEqual([{ index: 0, name: 'CT', isCT: true }]);
  });
  it('gray2d → one grayscale channel + .png; rgb2d → locked R,G,B + .png', () => {
    const g = applyDatasetKind(base().meta, 'gray2d');
    expect([g.dim, g.rgb, g.fileEnding]).toEqual(['2d', false, '.png']);
    expect(g.channels).toEqual([{ index: 0, name: 'grayscale', isCT: false }]);
    const r = applyDatasetKind(base().meta, 'rgb2d');
    expect([r.dim, r.rgb, r.fileEnding]).toEqual(['2d', true, '.png']);
    expect(r.channels.map((c) => c.name)).toEqual(['R', 'G', 'B']);
  });
});

describe('nnunetReducer — steps', () => {
  it('NEXT/BACK clamp to [0, last] and SET_STEP clamps', () => {
    let s = base();
    s = nnunetReducer(s, { type: 'BACK' });
    expect(s.step).toBe(0);
    for (let i = 0; i < 10; i++) s = nnunetReducer(s, { type: 'NEXT' });
    expect(s.step).toBe(NNUNET_STEPS.length - 1);
    expect(nnunetReducer(s, { type: 'SET_STEP', step: 99 }).step).toBe(NNUNET_STEPS.length - 1);
    expect(nnunetReducer(s, { type: 'SET_STEP', step: -5 }).step).toBe(0);
  });
});

describe('nnunetReducer — channels stay consecutive from 0', () => {
  it('ADD then REMOVE re-indexes', () => {
    let s = base();
    s = nnunetReducer(s, { type: 'ADD_CHANNEL' });
    s = nnunetReducer(s, { type: 'ADD_CHANNEL' });
    expect(s.meta.channels.map((c) => c.index)).toEqual([0, 1, 2]);
    s = nnunetReducer(s, { type: 'REMOVE_CHANNEL', index: 0 });
    expect(s.meta.channels.map((c) => c.index)).toEqual([0, 1]);
  });
  it('MOVE swaps and re-indexes; out-of-range is a no-op', () => {
    let s = base();
    s = nnunetReducer(s, { type: 'UPDATE_CHANNEL', index: 0, patch: { name: 'A' } });
    s = nnunetReducer(s, { type: 'ADD_CHANNEL' });
    s = nnunetReducer(s, { type: 'UPDATE_CHANNEL', index: 1, patch: { name: 'B' } });
    s = nnunetReducer(s, { type: 'MOVE_CHANNEL', index: 0, dir: 1 });
    expect(s.meta.channels.map((c) => `${c.index}:${c.name}`)).toEqual(['0:B', '1:A']);
    expect(nnunetReducer(s, { type: 'MOVE_CHANNEL', index: 0, dir: -1 })).toBe(s); // no-op returns same ref
  });
});

describe('nnunetReducer — labels stay consecutive + background protected', () => {
  it('REMOVE re-values; removing background (value 0) is a no-op', () => {
    let s = base();
    s = nnunetReducer(s, { type: 'ADD_LABEL' }); // bg0, label_1(1), ''(2)
    expect(s.meta.labels.map((l) => l.value)).toEqual([0, 1, 2]);
    s = nnunetReducer(s, { type: 'REMOVE_LABEL', index: 1 }); // drop label_1
    expect(s.meta.labels.map((l) => l.value)).toEqual([0, 1]);
    const before = s.meta.labels;
    s = nnunetReducer(s, { type: 'REMOVE_LABEL', index: 0 }); // background — protected
    expect(s.meta.labels.find((l) => l.value === 0)?.name).toBe('background');
    expect(s.meta.labels.length).toBe(before.length);
  });
});

describe('nnunetReducer — files + case split', () => {
  it('SET_CASE_SPLIT matches by inferred filename caseId; REMOVE_CASE drops the case; CLEAR empties', () => {
    let s = base();
    s = nnunetReducer(s, { type: 'ADD_FILES', files: [file({ filename: 'la_003_0000.nii.gz' }), file({ filename: 'la_003.nii.gz' }), file({ filename: 'la_004_0000.nii.gz' })] });
    expect(s.files).toHaveLength(3);
    s = nnunetReducer(s, { type: 'SET_CASE_SPLIT', caseId: 'la_003', split: 'test' });
    expect(s.files.filter((f) => f.filename.startsWith('la_003')).every((f) => f.split === 'test')).toBe(true);
    expect(s.files.find((f) => f.filename.startsWith('la_004'))?.split).toBe('train');
    s = nnunetReducer(s, { type: 'REMOVE_CASE', caseId: 'la_003' });
    expect(s.files.map((f) => f.filename)).toEqual(['la_004_0000.nii.gz']);
    expect(nnunetReducer(s, { type: 'CLEAR_FILES' }).files).toEqual([]);
  });
  it('SET_FILE_PNG attaches header info by id', () => {
    let s = base();
    const f = file();
    s = nnunetReducer(s, { type: 'ADD_FILES', files: [f] });
    s = nnunetReducer(s, { type: 'SET_FILE_PNG', id: f.id, png: { width: 256, height: 256, colorType: 0 } });
    expect(s.files[0].png).toEqual({ width: 256, height: 256, colorType: 0 });
  });
});
