import * as nifti from "nifti-reader-js";
import vtkImageData from "@kitware/vtk.js/Common/DataModel/ImageData";
import vtkDataArray from "@kitware/vtk.js/Common/Core/DataArray";
import type { NiftiData } from "./types";
import { binaryLabelMask, distinctLabelValues } from "./labelMask";

/**
 * Anatomical reorientation to closest-canonical RAS+ — a JS port of nibabel's
 * `as_closest_canonical`. The quad-view renderer assumes the voxel array is already in
 * canonical order (axis0 = L→R, axis1 = P→A, axis2 = I→S); TotalSegmentator output is
 * resampled to RAS so it renders correctly, but raw clinical data keeps its acquisition
 * orientation (e.g. KiTS/RibFrac files are `('I','P','L')`) and would otherwise render with
 * permuted/flipped axial/coronal/sagittal planes. We read the NIfTI affine and reorder + flip
 * the array to RAS+ so EVERY file matches the renderer's assumption.
 */
export type CanonicalOrientation = {
  /** perm[k] = the source array axis that becomes canonical axis k (0=R/X, 1=A/Y, 2=S/Z). */
  perm: [number, number, number];
  /** flips[k] = whether canonical axis k is reversed so +index points toward R / A / S. */
  flips: [boolean, boolean, boolean];
};

/**
 * Derive the RAS+ reorientation from a 4x4 NIfTI affine (`affine[worldRow][voxelCol]`, the
 * convention nifti-reader-js builds from sform/qform). Greedy axis assignment, mirroring
 * nibabel's `io_orientation`. Returns null when the affine is missing/singular so the caller
 * safely falls back to the raw array order.
 */
export function orientationFromAffine(
  affine: number[][] | undefined | null,
): CanonicalOrientation | null {
  if (!affine || affine.length < 3) return null;
  const m: number[][] = [];
  for (let w = 0; w < 3; w++) {
    const row = affine[w];
    if (!row || row.length < 3) return null;
    m.push([row[0], row[1], row[2]]);
  }
  const worldForInput = [-1, -1, -1];
  const signForInput = [1, 1, 1];
  const usedInput = [false, false, false];
  const usedWorld = [false, false, false];
  for (let pass = 0; pass < 3; pass++) {
    let best = -1;
    let bi = -1;
    let bw = -1;
    for (let w = 0; w < 3; w++) {
      if (usedWorld[w]) continue;
      for (let i = 0; i < 3; i++) {
        if (usedInput[i]) continue;
        const v = Math.abs(m[w][i]);
        if (v > best) {
          best = v;
          bw = w;
          bi = i;
        }
      }
    }
    if (bi < 0 || best === 0) return null; // singular / degenerate orientation
    worldForInput[bi] = bw;
    signForInput[bi] = m[bw][bi] < 0 ? -1 : 1;
    usedInput[bi] = true;
    usedWorld[bw] = true;
  }
  const perm: [number, number, number] = [0, 0, 0];
  const flips: [boolean, boolean, boolean] = [false, false, false];
  for (let i = 0; i < 3; i++) {
    const w = worldForInput[i];
    perm[w] = i;
    flips[w] = signForInput[i] < 0;
  }
  return { perm, flips };
}

/** True when the orientation is already RAS+ (no reorder, no flip) — the canonical fast path. */
export function isCanonicalOrientation(o: CanonicalOrientation): boolean {
  return (
    o.perm[0] === 0 &&
    o.perm[1] === 1 &&
    o.perm[2] === 2 &&
    !o.flips[0] &&
    !o.flips[1] &&
    !o.flips[2]
  );
}

/**
 * Reorder + flip a flat voxel array (axis0 fastest-varying) into RAS+ canonical order,
 * returning the canonical data + dims + spacing. O(N) with incremental strided indexing
 * (one read/write per voxel); only invoked for non-canonical files.
 */
export function applyCanonicalOrientation(
  data: Float32Array,
  dims: [number, number, number],
  spacing: [number, number, number],
  o: CanonicalOrientation,
): { data: Float32Array; dims: [number, number, number]; spacing: [number, number, number] } {
  const inDims = dims;
  const inStride = [1, inDims[0], inDims[0] * inDims[1]];
  const outDims: [number, number, number] = [dims[o.perm[0]], dims[o.perm[1]], dims[o.perm[2]]];
  const outSpacing: [number, number, number] = [
    spacing[o.perm[0]],
    spacing[o.perm[1]],
    spacing[o.perm[2]],
  ];
  const [m0, m1, m2] = outDims;
  // Per output axis k: the input-flat step for o_k+1, and the base offset (o_k=0).
  const step = [0, 0, 0];
  const start = [0, 0, 0];
  for (let k = 0; k < 3; k++) {
    const a = o.perm[k];
    step[k] = (o.flips[k] ? -1 : 1) * inStride[a];
    start[k] = o.flips[k] ? (inDims[a] - 1) * inStride[a] : 0;
  }
  const base = start[0] + start[1] + start[2];
  const out = new Float32Array(data.length);
  let outFlat = 0;
  let b2 = base;
  for (let o2 = 0; o2 < m2; o2++) {
    let b1 = b2;
    for (let o1 = 0; o1 < m1; o1++) {
      let inIdx = b1;
      for (let o0 = 0; o0 < m0; o0++) {
        out[outFlat++] = data[inIdx];
        inIdx += step[0];
      }
      b1 += step[1];
    }
    b2 += step[2];
  }
  return { data: out, dims: outDims, spacing: outSpacing };
}

/**
 * Parse a (possibly gzipped) NIfTI ArrayBuffer into a VTK ImageData + header/geometry.
 *
 * Pure + non-hook so it can be shared by BOTH the `useNiftiData` hook (the main image)
 * and the organ-mask overlay loader (which loads N extra volumes outside React state).
 * Throws on a non-NIfTI buffer. Mirrors the original in-hook conversion exactly
 * (datatype switch + scl_slope/scl_inter scaling + dims/spacing).
 */
export function parseNiftiBuffer(input: ArrayBuffer): NiftiData {
  let arrayBuffer: ArrayBuffer = input;
  if (nifti.isCompressed(arrayBuffer)) {
    arrayBuffer = nifti.decompress(arrayBuffer) as ArrayBuffer;
  }
  if (!nifti.isNIFTI(arrayBuffer)) {
    throw new Error("Not a valid NIfTI file");
  }

  const header = nifti.readHeader(arrayBuffer);
  if (!header) {
    throw new Error("Could not read NIfTI header");
  }
  const imageBuffer = nifti.readImage(header, arrayBuffer);
  if (!imageBuffer) {
    throw new Error("Could not read NIfTI image data");
  }

  const dims = header.dims;
  const nx = dims[1];
  const ny = dims[2];
  const nz = dims[3];

  const pixDims = header.pixDims;
  const sx = Math.abs(pixDims[1]) || 1;
  const sy = Math.abs(pixDims[2]) || 1;
  const sz = Math.abs(pixDims[3]) || 1;

  let typedData: Float32Array | Int16Array | Uint16Array | Uint8Array | Int8Array;
  switch (header.datatypeCode) {
    case nifti.NIFTI1.TYPE_UINT8:
      typedData = new Uint8Array(imageBuffer);
      break;
    case nifti.NIFTI1.TYPE_INT8:
      typedData = new Int8Array(imageBuffer);
      break;
    case nifti.NIFTI1.TYPE_UINT16:
      typedData = new Uint16Array(imageBuffer);
      break;
    case nifti.NIFTI1.TYPE_INT16:
      typedData = new Int16Array(imageBuffer);
      break;
    case nifti.NIFTI1.TYPE_FLOAT32:
      typedData = new Float32Array(imageBuffer);
      break;
    case nifti.NIFTI1.TYPE_FLOAT64: {
      const float64 = new Float64Array(imageBuffer);
      typedData = new Float32Array(float64);
      break;
    }
    case nifti.NIFTI1.TYPE_INT32: {
      const int32 = new Int32Array(imageBuffer);
      typedData = new Float32Array(int32);
      break;
    }
    default:
      typedData = new Float32Array(imageBuffer);
  }

  const slope = header.scl_slope || 1;
  const intercept = header.scl_inter || 0;
  let scaledData: Float32Array;
  if (slope !== 1 || intercept !== 0) {
    scaledData = new Float32Array(typedData.length);
    for (let i = 0; i < typedData.length; i++) {
      scaledData[i] = typedData[i] * slope + intercept;
    }
  } else {
    scaledData = typedData instanceof Float32Array ? typedData : new Float32Array(typedData);
  }

  // Reorient to closest-canonical RAS+ so a file in any acquisition orientation (e.g.
  // KiTS/RibFrac = ('I','P','L')) renders with the axial/coronal/sagittal mapping the
  // renderer assumes. Canonical files (TotalSegmentator = RAS) hit the no-op fast path.
  let outData: Float32Array = scaledData;
  let outDims: [number, number, number] = [nx, ny, nz];
  let outSpacing: [number, number, number] = [sx, sy, sz];
  const orientation = orientationFromAffine(header.affine);
  if (orientation && !isCanonicalOrientation(orientation)) {
    const reoriented = applyCanonicalOrientation(scaledData, [nx, ny, nz], [sx, sy, sz], orientation);
    outData = reoriented.data;
    outDims = reoriented.dims;
    outSpacing = reoriented.spacing;
  }

  const imageData = vtkImageData.newInstance();
  imageData.setDimensions(outDims);
  imageData.setSpacing(outSpacing);

  const dataArray = vtkDataArray.newInstance({
    name: "Scalars",
    numberOfComponents: 1,
    values: outData,
  });
  imageData.getPointData().setScalars(dataArray);

  return {
    header: {
      dims,
      pixDims,
      datatype: header.datatypeCode,
      littleEndian: header.littleEndian,
      voxOffset: header.vox_offset,
      affine: header.affine || [],
      description: header.description || "",
    },
    imageData,
    rawData: arrayBuffer,
    dimensions: outDims,
    spacing: outSpacing,
  };
}

/**
 * Load a NIfTI File into a VTK ImageData — for organ-mask overlays (`OrganMaskData.imageData`).
 * The mask must share the parent image's voxel grid (dims + spacing) to overlay correctly;
 * no resampling/affine alignment is applied here.
 */
export async function loadNiftiImageData(file: File): Promise<NiftiData["imageData"]> {
  const buffer = await file.arrayBuffer();
  return parseNiftiBuffer(buffer).imageData;
}

/** Distinct non-zero integer label values present in a loaded segmentation volume. */
export function labelValues(imageData: NiftiData["imageData"]): number[] {
  return distinctLabelValues(imageData.getPointData().getScalars().getData());
}

/**
 * Extract ONE organ from a multi-label volume as a binary (0/1) vtkImageData that shares the
 * source grid (dims / spacing / origin) — so the renderer's contour-at-0.5 + value-1 slice colour
 * pick out exactly that organ. Lets a single labelsTr/<case>.nii.gz drive N per-organ overlays.
 */
export function extractBinaryLabel(imageData: NiftiData["imageData"], value: number): NiftiData["imageData"] {
  const scalars = imageData.getPointData().getScalars().getData();
  const out = vtkImageData.newInstance();
  out.setDimensions(imageData.getDimensions());
  out.setSpacing(imageData.getSpacing());
  out.setOrigin(imageData.getOrigin());
  out.getPointData().setScalars(
    vtkDataArray.newInstance({ name: "Scalars", numberOfComponents: 1, values: binaryLabelMask(scalars, value) }),
  );
  return out;
}
