// Data-import DOMAIN — nnU-Net wizard state (pure reducer). Mirrors the feature's other FSMs:
// a pure (state, action) -> state function, illegal transitions are no-ops. The component holds
// this via useReducer and computes the derived views (groupFilesIntoCases / validateDataset /
// generateDatasetJson / planFolderTree from core.ts) with useMemo — they are NOT stored here.

import {
  endingsForDim,
  parseFilename,
  type Channel,
  type DatasetMeta,
  type FileEnding,
  type Label,
  type SourceFile,
  type Split,
} from './core';

export const NNUNET_STEPS = ['dataset', 'channels', 'labels', 'cases', 'export'] as const;
export type NnunetStepKey = (typeof NNUNET_STEPS)[number];

/** The three dataset shapes the wizard offers (drives dim/format/how channels map to files). */
export type DatasetKind = 'volume3d' | 'gray2d' | 'rgb2d';

export interface NnunetWizardState {
  /** 0..NNUNET_STEPS.length-1 */
  step: number;
  meta: DatasetMeta;
  /** Staged source files; role/caseId/channelIndex are inferred by core.groupFilesIntoCases. */
  files: SourceFile[];
}

export type NnunetAction =
  | { type: 'SET_STEP'; step: number }
  | { type: 'NEXT' }
  | { type: 'BACK' }
  | { type: 'SET_DATASET_KIND'; kind: DatasetKind }
  | { type: 'PATCH_META'; patch: Partial<Pick<DatasetMeta, 'id' | 'name' | 'fileEnding'>> }
  | { type: 'ADD_CHANNEL' }
  | { type: 'UPDATE_CHANNEL'; index: number; patch: Partial<Channel> }
  | { type: 'REMOVE_CHANNEL'; index: number }
  | { type: 'MOVE_CHANNEL'; index: number; dir: -1 | 1 }
  | { type: 'ADD_LABEL' }
  | { type: 'UPDATE_LABEL'; index: number; patch: Partial<Label> }
  | { type: 'REMOVE_LABEL'; index: number }
  | { type: 'ADD_FILES'; files: SourceFile[] }
  | { type: 'SET_FILE_PNG'; id: string; png: NonNullable<SourceFile['png']> }
  | { type: 'SET_CASE_SPLIT'; caseId: string; split: Split }
  | { type: 'REMOVE_CASE'; caseId: string }
  | { type: 'CLEAR_FILES' };

export const initialNnunetState: NnunetWizardState = {
  step: 0,
  meta: {
    id: 100,
    name: '',
    fileEnding: '.nii.gz',
    dim: '3d',
    rgb: false,
    channels: [{ index: 0, name: 'CT', isCT: true }],
    labels: [
      { name: 'background', value: 0 },
      { name: 'label_1', value: 1 },
    ],
  },
  files: [],
};

/** Re-index channels to 0..n-1 in array order (nnU-Net requires consecutive indices from 0). */
const reindexChannels = (channels: Channel[]): Channel[] => channels.map((c, index) => ({ ...c, index }));
/** Re-value labels to 0..n-1 in array order (background must stay 0, values consecutive). */
const revalueLabels = (labels: Label[]): Label[] => labels.map((l, value) => ({ ...l, value }));

/** Apply a dataset-kind preset: sets dim/rgb/fileEnding and a sensible default channel set. */
export function applyDatasetKind(meta: DatasetMeta, kind: DatasetKind): DatasetMeta {
  if (kind === 'volume3d') {
    const wasRgb = meta.channels.map((c) => c.name.toUpperCase()).join(',') === 'R,G,B';
    const keepEnding = (['.nii.gz', '.nrrd', '.mha'] as FileEnding[]).includes(meta.fileEnding);
    return {
      ...meta,
      dim: '3d',
      rgb: false,
      fileEnding: keepEnding ? meta.fileEnding : '.nii.gz',
      channels: wasRgb ? [{ index: 0, name: 'CT', isCT: true }] : meta.channels,
    };
  }
  if (kind === 'gray2d') {
    return { ...meta, dim: '2d', rgb: false, fileEnding: '.png', channels: [{ index: 0, name: 'grayscale', isCT: false }] };
  }
  // rgb2d
  return {
    ...meta,
    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 },
    ],
  };
}

/** The case id a staged file belongs to (explicit, else inferred from its nnU-Net filename). */
const caseIdOf = (f: SourceFile, ending: FileEnding): string | undefined =>
  f.caseId ?? parseFilename(f.filename, ending)?.caseId;

export function nnunetReducer(state: NnunetWizardState, action: NnunetAction): NnunetWizardState {
  switch (action.type) {
    case 'SET_STEP':
      return { ...state, step: clampStep(action.step) };
    case 'NEXT':
      return { ...state, step: clampStep(state.step + 1) };
    case 'BACK':
      return { ...state, step: clampStep(state.step - 1) };

    case 'SET_DATASET_KIND':
      return { ...state, meta: applyDatasetKind(state.meta, action.kind) };
    case 'PATCH_META':
      return { ...state, meta: { ...state.meta, ...action.patch } };

    case 'ADD_CHANNEL':
      return {
        ...state,
        meta: { ...state.meta, channels: [...state.meta.channels, { index: state.meta.channels.length, name: '', isCT: false }] },
      };
    case 'UPDATE_CHANNEL':
      return {
        ...state,
        meta: { ...state.meta, channels: state.meta.channels.map((c, i) => (i === action.index ? { ...c, ...action.patch } : c)) },
      };
    case 'REMOVE_CHANNEL':
      return {
        ...state,
        meta: { ...state.meta, channels: reindexChannels(state.meta.channels.filter((_, i) => i !== action.index)) },
      };
    case 'MOVE_CHANNEL': {
      const j = action.index + action.dir;
      if (j < 0 || j >= state.meta.channels.length) return state;
      const channels = [...state.meta.channels];
      [channels[action.index], channels[j]] = [channels[j], channels[action.index]];
      return { ...state, meta: { ...state.meta, channels: reindexChannels(channels) } };
    }

    case 'ADD_LABEL':
      return { ...state, meta: { ...state.meta, labels: [...state.meta.labels, { name: '', value: state.meta.labels.length }] } };
    case 'UPDATE_LABEL':
      return {
        ...state,
        meta: { ...state.meta, labels: state.meta.labels.map((l, i) => (i === action.index ? { ...l, ...action.patch } : l)) },
      };
    case 'REMOVE_LABEL':
      // Never remove background (value 0); always re-value the rest consecutively.
      return {
        ...state,
        meta: { ...state.meta, labels: revalueLabels(state.meta.labels.filter((l, i) => i !== action.index || l.value === 0)) },
      };

    case 'ADD_FILES':
      return { ...state, files: [...state.files, ...action.files] };
    case 'SET_FILE_PNG':
      return { ...state, files: state.files.map((f) => (f.id === action.id ? { ...f, png: action.png } : f)) };
    case 'SET_CASE_SPLIT':
      return {
        ...state,
        files: state.files.map((f) => (caseIdOf(f, state.meta.fileEnding) === action.caseId ? { ...f, split: action.split } : f)),
      };
    case 'REMOVE_CASE':
      return { ...state, files: state.files.filter((f) => caseIdOf(f, state.meta.fileEnding) !== action.caseId) };
    case 'CLEAR_FILES':
      return { ...state, files: [] };

    default:
      return state;
  }
}

function clampStep(n: number): number {
  return Math.max(0, Math.min(NNUNET_STEPS.length - 1, n));
}

/** Accepted file endings for the current dim — for the format <select> in the Dataset step. */
export const endingsForState = (state: NnunetWizardState): FileEnding[] => endingsForDim(state.meta.dim);
