import { describe, it, expect } from 'vitest';
import {
  pad3,
  pad4,
  datasetFolderName,
  parseFilename,
  groupFilesIntoCases,
  validateDataset,
  generateDatasetJson,
  planFolderTree,
  generateBashScript,
  generatePythonScript,
  type Channel,
  type Label,
  type DatasetMeta,
  type DatasetCase,
  type SourceFile,
  type FileEnding,
  type Split,
  type ExecOptions,
} from './core';

// ----------------------------- fixtures -----------------------------

/** A valid DatasetMeta; override any field. Default: 3D, single non-CT channel, bg + 1 label. */
function mkMeta(over: Partial<DatasetMeta> = {}): DatasetMeta {
  return {
    id: 5,
    name: 'Prostate',
    fileEnding: '.nii.gz',
    dim: '3d',
    rgb: false,
    channels: [{ index: 0, name: 'T2', isCT: false }],
    labels: [
      { name: 'background', value: 0 },
      { name: 'PZ', value: 1 },
    ],
    ...over,
  };
}

/** A train case with channel 0 + a label, named to match meta.fileEnding. */
function mkCase(over: Partial<DatasetCase> = {}, ending: FileEnding = '.nii.gz'): DatasetCase {
  const caseId = over.caseId ?? 'BRATS_001';
  const split: Split = over.split ?? 'train';
  const channelFiles =
    over.channelFiles ?? {
      0: mkFile(`${caseId}_0000${ending}`, { role: 'image', caseId, channelIndex: 0, split }),
    };
  const labelFile =
    'labelFile' in over
      ? over.labelFile
      : mkFile(`${caseId}${ending}`, { role: 'label', caseId, split });
  return { caseId, split, channelFiles, labelFile };
}

function mkFile(filename: string, over: Partial<SourceFile> = {}): SourceFile {
  return {
    id: over.id ?? `f-${filename}`,
    filename,
    sizeBytes: over.sizeBytes ?? 100,
    role: over.role ?? 'unknown',
    caseId: over.caseId,
    channelIndex: over.channelIndex,
    split: over.split ?? 'train',
    png: over.png,
    blob: over.blob,
  };
}

const codes = (r: { issues: { code: string }[] }) => r.issues.map((i) => i.code);
const defOpts = (over: Partial<ExecOptions> = {}): ExecOptions => ({
  operation: 'copy',
  target: 'local',
  sourceRoot: '/in',
  destRoot: '/out',
  ...over,
});

// =============================== padding ===============================

describe('pad3 / pad4', () => {
  it('zero-pads to width 3 and 4', () => {
    expect(pad3(5)).toBe('005');
    expect(pad3(42)).toBe('042');
    expect(pad3(999)).toBe('999');
    expect(pad4(0)).toBe('0000');
    expect(pad4(12)).toBe('0012');
  });
});

// =========================== datasetFolderName ===========================

describe('datasetFolderName', () => {
  it('pads id to 3 and strips spaces from the name', () => {
    expect(datasetFolderName(mkMeta({ id: 5, name: 'Prostate' }))).toBe('Dataset005_Prostate');
    expect(datasetFolderName(mkMeta({ id: 5, name: 'Pro state' }))).toBe('Dataset005_Prostate');
    expect(datasetFolderName(mkMeta({ id: 123, name: 'Brain Tumor Seg' }))).toBe(
      'Dataset123_BrainTumorSeg',
    );
  });
});

// ============================= parseFilename =============================

describe('parseFilename', () => {
  it('parses an image filename with a 4-digit channel suffix', () => {
    expect(parseFilename('BRATS_001_0000.nii.gz', '.nii.gz')).toEqual({
      caseId: 'BRATS_001',
      channelIndex: 0,
      role: 'image',
    });
    expect(parseFilename('BRATS_001_0003.nii.gz', '.nii.gz')).toEqual({
      caseId: 'BRATS_001',
      channelIndex: 3,
      role: 'image',
    });
  });

  it('parses a label filename (no channel suffix)', () => {
    expect(parseFilename('BRATS_001.nii.gz', '.nii.gz')).toEqual({
      caseId: 'BRATS_001',
      role: 'label',
    });
  });

  it('returns null when the ending does not match', () => {
    expect(parseFilename('BRATS_001_0000.nrrd', '.nii.gz')).toBeNull();
    expect(parseFilename('BRATS_001.png', '.nii.gz')).toBeNull();
  });
});

// ========================== groupFilesIntoCases ==========================

describe('groupFilesIntoCases', () => {
  it('groups a flat file list by caseId into channelFiles + labelFile', () => {
    const meta = mkMeta();
    const files: SourceFile[] = [
      mkFile('BRATS_001_0000.nii.gz'),
      mkFile('BRATS_001.nii.gz'),
      mkFile('BRATS_002_0000.nii.gz'),
    ];
    const { cases, unassigned } = groupFilesIntoCases(files, meta);
    expect(unassigned).toHaveLength(0);
    expect(cases).toHaveLength(2);

    const c1 = cases.find((c) => c.caseId === 'BRATS_001')!;
    expect(c1.channelFiles[0]?.filename).toBe('BRATS_001_0000.nii.gz');
    expect(c1.channelFiles[0]?.role).toBe('image');
    expect(c1.channelFiles[0]?.channelIndex).toBe(0);
    expect(c1.labelFile?.filename).toBe('BRATS_001.nii.gz');
    expect(c1.labelFile?.role).toBe('label');

    const c2 = cases.find((c) => c.caseId === 'BRATS_002')!;
    expect(c2.channelFiles[0]?.filename).toBe('BRATS_002_0000.nii.gz');
    expect(c2.labelFile).toBeUndefined();
  });

  it('puts unparseable files (wrong ending, no caseId) into the unassigned bucket', () => {
    const meta = mkMeta(); // .nii.gz
    const files: SourceFile[] = [
      mkFile('BRATS_001_0000.nii.gz'),
      mkFile('stray.txt'), // does not end with .nii.gz, role unknown -> unassigned
    ];
    const { cases, unassigned } = groupFilesIntoCases(files, meta);
    expect(cases).toHaveLength(1);
    expect(unassigned).toHaveLength(1);
    expect(unassigned[0].filename).toBe('stray.txt');
  });

  it('honours a pre-set caseId/role instead of re-parsing the filename', () => {
    const meta = mkMeta();
    const files: SourceFile[] = [
      mkFile('weird-name.nii.gz', { role: 'image', caseId: 'CASE_X', channelIndex: 0 }),
    ];
    const { cases } = groupFilesIntoCases(files, meta);
    expect(cases).toHaveLength(1);
    expect(cases[0].caseId).toBe('CASE_X');
    expect(cases[0].channelFiles[0]?.filename).toBe('weird-name.nii.gz');
  });
});

// ============================ validateDataset ============================

describe('validateDataset — valid input', () => {
  it('a well-formed single-channel 3D dataset yields ok:true with 0 errors', () => {
    // id 5 -> dataset.idReserved warning is expected, so use id 42 to get a fully clean report.
    const meta = mkMeta({ id: 42, name: 'Prostate' });
    const r = validateDataset(meta, [mkCase()]);
    expect(r.ok).toBe(true);
    expect(r.errorCount).toBe(0);
    expect(r.warningCount).toBe(0);
    expect(r.issues).toHaveLength(0);
  });

  it('bookkeeping: ok=false once any error is present; counts split by level', () => {
    const meta = mkMeta({ id: 5, name: '' }); // id 5 -> reserved warning; empty name -> error
    const r = validateDataset(meta, [mkCase()]);
    expect(r.ok).toBe(false);
    expect(r.errorCount).toBeGreaterThanOrEqual(1);
    expect(r.warningCount).toBeGreaterThanOrEqual(1);
    expect(r.errorCount + r.warningCount).toBe(r.issues.length);
    expect(codes(r)).toContain('dataset.name');
    expect(codes(r)).toContain('dataset.idReserved');
  });
});

describe('validateDataset — dataset id / name rules', () => {
  it('dataset.id: id below 1 or above 999 is an error', () => {
    expect(codes(validateDataset(mkMeta({ id: 0 }), [mkCase()]))).toContain('dataset.id');
    expect(codes(validateDataset(mkMeta({ id: 1000 }), [mkCase()]))).toContain('dataset.id');
  });

  it('dataset.idReserved: ids 1..10 emit a warning', () => {
    const r = validateDataset(mkMeta({ id: 3 }), [mkCase()]);
    const reserved = r.issues.find((i) => i.code === 'dataset.idReserved');
    expect(reserved?.level).toBe('warning');
  });

  it('dataset.name: empty / whitespace name is an error', () => {
    expect(codes(validateDataset(mkMeta({ id: 42, name: '   ' }), [mkCase()]))).toContain(
      'dataset.name',
    );
  });

  it('dataset.nameSpaces: a name with spaces is a warning', () => {
    const r = validateDataset(mkMeta({ id: 42, name: 'Has Space' }), [mkCase()]);
    const sp = r.issues.find((i) => i.code === 'dataset.nameSpaces');
    expect(sp?.level).toBe('warning');
  });
});

describe('validateDataset — channel rules', () => {
  it('channels.empty: no channels is an error', () => {
    const meta = mkMeta({ id: 42, channels: [] });
    // case must satisfy whatever fileChannels yields; with no channels there is no required channel.
    const c = mkCase({ channelFiles: {} });
    expect(codes(validateDataset(meta, [c]))).toContain('channels.empty');
  });

  it('channels.consecutive: a non-0-based / gapped index set is an error', () => {
    const meta = mkMeta({
      id: 42,
      channels: [{ index: 1, name: 'T2', isCT: false }], // starts at 1, not 0
    });
    const c = mkCase({
      channelFiles: { 1: mkFile('BRATS_001_0001.nii.gz', { role: 'image', caseId: 'BRATS_001', channelIndex: 1 }) },
    });
    expect(codes(validateDataset(meta, [c]))).toContain('channels.consecutive');
  });

  it('channels.name: a blank channel name is an error', () => {
    const meta = mkMeta({ id: 42, channels: [{ index: 0, name: '  ', isCT: false }] });
    expect(codes(validateDataset(meta, [mkCase()]))).toContain('channels.name');
  });

  it('channels.dupName: duplicate channel names are a warning', () => {
    const meta = mkMeta({
      id: 42,
      channels: [
        { index: 0, name: 'T2', isCT: false },
        { index: 1, name: 'T2', isCT: false },
      ],
    });
    const c = mkCase({
      channelFiles: {
        0: mkFile('BRATS_001_0000.nii.gz', { role: 'image', caseId: 'BRATS_001', channelIndex: 0 }),
        1: mkFile('BRATS_001_0001.nii.gz', { role: 'image', caseId: 'BRATS_001', channelIndex: 1 }),
      },
    });
    const r = validateDataset(meta, [c]);
    const dup = r.issues.find((i) => i.code === 'channels.dupName');
    expect(dup?.level).toBe('warning');
  });
});

describe('validateDataset — label rules', () => {
  it('labels.background: missing value-0 label is an error', () => {
    const meta = mkMeta({ id: 42, labels: [{ name: 'PZ', value: 1 }] });
    expect(codes(validateDataset(meta, [mkCase()]))).toContain('labels.background');
  });

  it('labels.consecutive: gapped label values are an error', () => {
    const meta = mkMeta({
      id: 42,
      labels: [
        { name: 'background', value: 0 },
        { name: 'PZ', value: 2 }, // gap: no value 1
      ],
    });
    expect(codes(validateDataset(meta, [mkCase()]))).toContain('labels.consecutive');
  });

  it('labels.name: a blank label name is an error', () => {
    const meta = mkMeta({
      id: 42,
      labels: [
        { name: 'background', value: 0 },
        { name: '   ', value: 1 },
      ],
    });
    expect(codes(validateDataset(meta, [mkCase()]))).toContain('labels.name');
  });

  it('labels.dupName: duplicate label names are an error', () => {
    const meta = mkMeta({
      id: 42,
      labels: [
        { name: 'same', value: 0 },
        { name: 'same', value: 1 },
      ],
    });
    const r = validateDataset(meta, [mkCase()]);
    const dup = r.issues.find((i) => i.code === 'labels.dupName');
    expect(dup?.level).toBe('error');
  });
});

describe('validateDataset — format / dimension coherence', () => {
  it('format.dimMismatch: a 2D-only ending on a 3D dataset is a warning', () => {
    // .png is a 2D ending; on a 3d dataset this is unusual -> warning.
    const meta = mkMeta({ id: 42, dim: '3d', fileEnding: '.png' });
    const c = mkCase({}, '.png');
    const r = validateDataset(meta, [c]);
    const mm = r.issues.find((i) => i.code === 'format.dimMismatch');
    expect(mm?.level).toBe('warning');
  });

  it('rgb.channels (warning) + rgb.format (error): RGB on a non-image ending with wrong channels', () => {
    // rgb true, .nii.gz (not a 2D image format) and a single non-R/G/B channel.
    const meta = mkMeta({
      id: 42,
      dim: '2d',
      rgb: true,
      fileEnding: '.nii.gz',
      channels: [{ index: 0, name: 'T2', isCT: false }],
    });
    // RGB => single file slot at channel 0.
    const c = mkCase({}, '.nii.gz');
    const r = validateDataset(meta, [c]);
    const rgbCh = r.issues.find((i) => i.code === 'rgb.channels');
    const rgbFmt = r.issues.find((i) => i.code === 'rgb.format');
    expect(rgbCh?.level).toBe('warning');
    expect(rgbFmt?.level).toBe('error');
  });

  it('a correct RGB dataset (3x R/G/B, .png) emits neither rgb.channels nor rgb.format', () => {
    const meta = mkMeta({
      id: 42,
      dim: '2d',
      rgb: true,
      fileEnding: '.png',
      channels: [
        { index: 0, name: 'R', isCT: false },
        { index: 1, name: 'G', isCT: false },
        { index: 2, name: 'B', isCT: false },
      ],
    });
    const c = mkCase({}, '.png'); // RGB -> only channel 0 slot required
    const r = validateDataset(meta, [c]);
    expect(codes(r)).not.toContain('rgb.channels');
    expect(codes(r)).not.toContain('rgb.format');
  });
});

describe('validateDataset — case rules', () => {
  it('cases.noTrain: a dataset with only test cases is an error', () => {
    const meta = mkMeta({ id: 42 });
    const testCase = mkCase({ caseId: 'T1', split: 'test', labelFile: undefined });
    expect(codes(validateDataset(meta, [testCase]))).toContain('cases.noTrain');
  });

  it('cases.dupId: two cases with the same id is an error', () => {
    const meta = mkMeta({ id: 42 });
    const a = mkCase({ caseId: 'BRATS_001' });
    const b = mkCase({ caseId: 'BRATS_001' });
    expect(codes(validateDataset(meta, [a, b]))).toContain('cases.dupId');
  });

  it('case.missingChannel: a case lacking a required channel file is an error scoped to the case', () => {
    const meta = mkMeta({ id: 42 });
    const c = mkCase({ caseId: 'BRATS_001', channelFiles: {} }); // missing channel 0
    const r = validateDataset(meta, [c]);
    const miss = r.issues.find((i) => i.code === 'case.missingChannel');
    expect(miss).toBeDefined();
    expect(miss?.caseId).toBe('BRATS_001');
  });

  it('case.missingLabel: a train case with no label file is an error', () => {
    const meta = mkMeta({ id: 42 });
    const c = mkCase({ caseId: 'BRATS_001', labelFile: undefined });
    const r = validateDataset(meta, [c]);
    const miss = r.issues.find((i) => i.code === 'case.missingLabel');
    expect(miss?.level).toBe('error');
    expect(miss?.caseId).toBe('BRATS_001');
  });

  it('case.testLabel: a test case that carries a label is a warning', () => {
    const meta = mkMeta({ id: 42 });
    const train = mkCase({ caseId: 'TR1' });
    const test = mkCase({
      caseId: 'TS1',
      split: 'test',
      channelFiles: {
        0: mkFile('TS1_0000.nii.gz', { role: 'image', caseId: 'TS1', channelIndex: 0, split: 'test' }),
      },
      labelFile: mkFile('TS1.nii.gz', { role: 'label', caseId: 'TS1', split: 'test' }),
    });
    const r = validateDataset(meta, [train, test]);
    const tl = r.issues.find((i) => i.code === 'case.testLabel');
    expect(tl?.level).toBe('warning');
    expect(tl?.caseId).toBe('TS1');
  });
});

describe('validateDataset — file format guards', () => {
  it('file.lossy: a .jpg file is an error', () => {
    const meta = mkMeta({ id: 42, dim: '2d', fileEnding: '.png' });
    const c = mkCase(
      {
        caseId: 'C1',
        channelFiles: {
          0: mkFile('C1_0000.jpg', { role: 'image', caseId: 'C1', channelIndex: 0 }),
        },
        labelFile: mkFile('C1.png', { role: 'label', caseId: 'C1' }),
      },
      '.png',
    );
    expect(codes(validateDataset(meta, [c]))).toContain('file.lossy');
  });

  it('file.ending: a file not ending in the declared ending is an error', () => {
    const meta = mkMeta({ id: 42, fileEnding: '.nii.gz' });
    const c = mkCase({
      caseId: 'C1',
      channelFiles: {
        0: mkFile('C1_0000.nrrd', { role: 'image', caseId: 'C1', channelIndex: 0 }),
      },
      labelFile: mkFile('C1.nii.gz', { role: 'label', caseId: 'C1' }),
    });
    expect(codes(validateDataset(meta, [c]))).toContain('file.ending');
  });
});

// =========================== generateDatasetJson ===========================

describe('generateDatasetJson', () => {
  it('keys channel_names by index, mapping isCT channels to "CT"', () => {
    const meta = mkMeta({
      channels: [
        { index: 0, name: 'scan', isCT: true },
        { index: 1, name: 'T2', isCT: false },
      ],
    });
    const json = generateDatasetJson(meta, [mkCase()]) as {
      channel_names: Record<string, string>;
    };
    expect(json.channel_names['0']).toBe('CT');
    expect(json.channel_names['1']).toBe('T2');
  });

  it('maps labels name->value, counts numTraining, and echoes file_ending', () => {
    const meta = mkMeta({
      labels: [
        { name: 'background', value: 0 },
        { name: 'PZ', value: 1 },
        { name: 'TZ', value: 2 },
      ],
    });
    const cases: DatasetCase[] = [
      mkCase({ caseId: 'A' }),
      mkCase({ caseId: 'B' }),
      mkCase({
        caseId: 'T1',
        split: 'test',
        channelFiles: {
          0: mkFile('T1_0000.nii.gz', { role: 'image', caseId: 'T1', channelIndex: 0, split: 'test' }),
        },
        labelFile: undefined,
      }),
    ];
    const json = generateDatasetJson(meta, cases) as {
      labels: Record<string, number>;
      numTraining: number;
      file_ending: string;
    };
    expect(json.labels).toEqual({ background: 0, PZ: 1, TZ: 2 });
    expect(json.numTraining).toBe(2); // only train cases A + B
    expect(json.file_ending).toBe('.nii.gz');
  });
});

// ============================= planFolderTree =============================

describe('planFolderTree', () => {
  it('plans a train case into imagesTr/<case>_<XXXX> and labelsTr/<case>', () => {
    const meta = mkMeta({ id: 5, name: 'Prostate' });
    const plan = planFolderTree(meta, [mkCase({ caseId: 'BRATS_001' })]);
    const targets = plan.map((p) => p.targetPath);
    expect(targets).toContain('Dataset005_Prostate/imagesTr/BRATS_001_0000.nii.gz');
    expect(targets).toContain('Dataset005_Prostate/labelsTr/BRATS_001.nii.gz');
  });

  it('plans a test case into imagesTs and emits no label path', () => {
    const meta = mkMeta({ id: 5, name: 'Prostate' });
    const test = mkCase({
      caseId: 'TS1',
      split: 'test',
      channelFiles: {
        0: mkFile('TS1_0000.nii.gz', { role: 'image', caseId: 'TS1', channelIndex: 0, split: 'test' }),
      },
      labelFile: mkFile('TS1.nii.gz', { role: 'label', caseId: 'TS1', split: 'test' }),
    });
    const plan = planFolderTree(meta, [test]);
    const targets = plan.map((p) => p.targetPath);
    expect(targets).toContain('Dataset005_Prostate/imagesTs/TS1_0000.nii.gz');
    expect(targets.some((t) => t.includes('labelsTr') || t.includes('labelsTs'))).toBe(false);
  });

  it('RGB uses a single file slot (channel 0) regardless of declared channels', () => {
    const meta = mkMeta({
      id: 5,
      name: 'Cells',
      dim: '2d',
      rgb: true,
      fileEnding: '.png',
      channels: [
        { index: 0, name: 'R', isCT: false },
        { index: 1, name: 'G', isCT: false },
        { index: 2, name: 'B', isCT: false },
      ],
    });
    const c = mkCase({ caseId: 'CELL1' }, '.png');
    const plan = planFolderTree(meta, [c]);
    const imageTargets = plan.filter((p) => p.targetPath.includes('imagesTr'));
    expect(imageTargets).toHaveLength(1); // only channel 0 slot
    expect(imageTargets[0].targetPath).toBe('Dataset005_Cells/imagesTr/CELL1_0000.png');
  });
});

// ===================== generateBashScript / PythonScript =====================

describe('generateBashScript', () => {
  it('embeds the dataset.json block, mkdir of the dataset folder, and source->target manifest lines', () => {
    const meta = mkMeta({ id: 5, name: 'Prostate' });
    const cases = [mkCase({ caseId: 'BRATS_001' })];
    const plan = planFolderTree(meta, cases);
    const json = generateDatasetJson(meta, cases);
    const script = generateBashScript(meta, plan, json, defOpts());

    // dataset.json block present (serialized JSON contains the file_ending key)
    expect(script).toContain('dataset.json');
    expect(script).toContain('"file_ending": ".nii.gz"');
    // mkdir of a path under the dataset folder
    expect(script).toContain('mkdir -p "$DEST_ROOT/Dataset005_Prostate/imagesTr"');
    // manifest: src<TAB>dst line for the channel file
    expect(script).toContain('BRATS_001_0000.nii.gz\tDataset005_Prostate/imagesTr/BRATS_001_0000.nii.gz');
    // honours the copy operation
    expect(script).toContain('cp -p');
  });

  it('uses mv when operation is move', () => {
    const meta = mkMeta({ id: 5, name: 'Prostate' });
    const cases = [mkCase({ caseId: 'BRATS_001' })];
    const plan = planFolderTree(meta, cases);
    const json = generateDatasetJson(meta, cases);
    const script = generateBashScript(meta, plan, json, defOpts({ operation: 'move' }));
    expect(script).toContain('mv "$SOURCE_DIR/$src" "$DEST_ROOT/$dst"');
  });
});

describe('generatePythonScript', () => {
  it('embeds DATASET_JSON, makedirs of the dataset folder, and source->target manifest tuples', () => {
    const meta = mkMeta({ id: 5, name: 'Prostate' });
    const cases = [mkCase({ caseId: 'BRATS_001' })];
    const plan = planFolderTree(meta, cases);
    const json = generateDatasetJson(meta, cases);
    const script = generatePythonScript(meta, plan, json, defOpts());

    expect(script).toContain('DATASET_JSON = ');
    expect(script).toContain('"file_ending": ".nii.gz"');
    // makedirs of the dataset root + subdirs
    expect(script).toContain('os.path.join(DEST_ROOT, "Dataset005_Prostate")');
    expect(script).toContain('os.makedirs(os.path.join(root, d), exist_ok=True)');
    // manifest tuple: (src, dst)
    expect(script).toContain(
      '("BRATS_001_0000.nii.gz", "Dataset005_Prostate/imagesTr/BRATS_001_0000.nii.gz"),',
    );
  });
});
