/**
 * nnU-Net v2 dataset organizer — framework-agnostic typed core.
 *
 * Combines nnU-Net v2 dataset-format requirements with the RedBrick import model:
 *   RedBrick Task   -> nnU-Net training/test CASE
 *   RedBrick series -> nnU-Net input CHANNEL (modality)
 *   RedBrick GT     -> nnU-Net LABEL (segmentation in labelsTr)
 *
 * Pure functions only — no React, no DOM. Import into UI or a build step.
 *
 * Ported verbatim into the data-import feature from docs/workflows/nnunet.ts (the design
 * reference). This is the pure domain for the nnU-Net upload mode; the UI wraps it. Keep its
 * lossless-only FILE_ENDINGS separate from the feature's input formats.ts (opposite .jpg policy).
 */

// ----------------------------- Types -----------------------------

/** Lossless formats nnU-Net supports. .jpg/.jpeg are intentionally excluded. */
export type FileEnding =
  | ".nii.gz"
  | ".nrrd"
  | ".mha"
  | ".png"
  | ".bmp"
  | ".tif"
  | ".tiff";

export const FILE_ENDINGS: FileEnding[] = [
  ".nii.gz",
  ".nrrd",
  ".mha",
  ".png",
  ".bmp",
  ".tif",
  ".tiff",
];

export type Split = "train" | "test";

export interface Channel {
  /** 0-based, must be consecutive across the dataset. Rendered as 4-digit XXXX. */
  index: number;
  /** e.g. "T2", "ADC", "CT". "CT" triggers global intensity normalization. */
  name: string;
  isCT: boolean;
}

export interface Label {
  /** e.g. "background", "PZ". */
  name: string;
  /** Integer class value. Must be consecutive from 0; background must be 0. */
  value: number;
}

export interface PngInfo {
  width: number;
  height: number;
  /** PNG IHDR colour type: 0 gray, 2 RGB, 3 indexed, 4 gray+alpha, 6 RGBA. */
  colorType: number;
}

export interface SourceFile {
  id: string;
  filename: string;
  sizeBytes: number;
  /** Inferred or user-set: image (has _XXXX) vs label (no channel suffix). */
  role: "image" | "label" | "unknown";
  caseId?: string;
  channelIndex?: number;
  split: Split;
  /** Retained handle (enables PNG inspection and presigned upload). */
  blob?: Blob;
  /** Parsed from the PNG header when available. */
  png?: PngInfo;
}

export interface DatasetCase {
  caseId: string;
  split: Split;
  /** channelIndex -> file (undefined = missing channel). */
  channelFiles: Record<number, SourceFile | undefined>;
  /** Required for train cases, absent for test cases. */
  labelFile?: SourceFile;
}

export interface DatasetMeta {
  /** 1..999, rendered 3-digit. 1..10 reserved for the Medical Segmentation Decathlon. */
  id: number;
  /** Free-form, conventionally CamelCase with no spaces. */
  name: string;
  fileEnding: FileEnding;
  /** Spatial dimensionality. Drives default formats and the training config (2d vs 3d_fullres). */
  dim: "2d" | "3d";
  /** RGB natural-image mode: all three colour channels live in ONE file (nnU-Net's exception). */
  rgb: boolean;
  channels: Channel[];
  labels: Label[];
}

export type IssueLevel = "error" | "warning";

export interface ValidationIssue {
  level: IssueLevel;
  /** Stable rule id, e.g. "labels.consecutive". */
  code: string;
  message: string;
  /** Optional case the issue is scoped to. */
  caseId?: string;
}

export interface ValidationReport {
  issues: ValidationIssue[];
  errorCount: number;
  warningCount: number;
  ok: boolean;
}

export interface PlannedFile {
  source: SourceFile;
  /** Path relative to nnUNet_raw, e.g. Dataset005_Prostate/imagesTr/case_0000.nii.gz */
  targetPath: string;
}

// --------------------------- Formatting ---------------------------

export const pad3 = (n: number): string => String(n).padStart(3, "0");
export const pad4 = (n: number): string => String(n).padStart(4, "0");

/** "Dataset005_Prostate" */
export function datasetFolderName(meta: DatasetMeta): string {
  const safeName = (meta.name || "Dataset").replace(/\s+/g, "");
  return `Dataset${pad3(meta.id)}_${safeName}`;
}

/** 2D natural-image formats vs 3D volume formats. */
export function endingsForDim(dim: "2d" | "3d"): FileEnding[] {
  return dim === "2d" ? [".png", ".bmp", ".tif", ".tiff"] : [".nii.gz", ".nrrd", ".mha", ".tif", ".tiff"];
}

/**
 * Files-per-case channel list. For RGB the three colour channels live in ONE
 * file, so there is a single file slot (channel 0) even though dataset.json
 * still declares three channel_names (R, G, B).
 */
export function fileChannels(meta: DatasetMeta): { index: number; name: string }[] {
  return meta.rgb ? [{ index: 0, name: "RGB" }] : meta.channels.map((c) => ({ index: c.index, name: c.name }));
}

/** Parse width/height/colour-type from a PNG's first bytes (IHDR). Null if not a PNG. */
export function parsePngHeader(bytes: Uint8Array): PngInfo | null {
  if (bytes.length < 26) return null;
  const sig = [137, 80, 78, 71, 13, 10, 26, 10];
  for (let i = 0; i < 8; i++) if (bytes[i] !== sig[i]) return null;
  const u32 = (o: number) => ((bytes[o] << 24) | (bytes[o + 1] << 16) | (bytes[o + 2] << 8) | bytes[o + 3]) >>> 0;
  return { width: u32(16), height: u32(20), colorType: bytes[25] };
}

export function pngColorTypeName(ct: number): string {
  return ({ 0: "grayscale", 2: "RGB", 3: "indexed", 4: "gray+alpha", 6: "RGBA" } as Record<number, string>)[ct] ?? `type ${ct}`;
}

// ------------------------- Filename parsing -------------------------

/** Strip a known file ending (handles multi-dot like .nii.gz). */
export function stripEnding(filename: string, ending: FileEnding): string | null {
  return filename.endsWith(ending) ? filename.slice(0, -ending.length) : null;
}

/**
 * Infer {caseId, channelIndex, role} from an nnU-Net-style filename.
 * Images end with _XXXX (4 digits); labels do not.
 *   "BRATS_001_0000.nii.gz" -> { caseId: "BRATS_001", channelIndex: 0, role: "image" }
 *   "BRATS_001.nii.gz"      -> { caseId: "BRATS_001", role: "label" }
 */
export function parseFilename(
  filename: string,
  ending: FileEnding
): { caseId: string; channelIndex?: number; role: "image" | "label" } | null {
  const base = stripEnding(filename, ending);
  if (base === null) return null;
  const m = base.match(/^(.*)_(\d{4})$/);
  if (m) return { caseId: m[1], channelIndex: parseInt(m[2], 10), role: "image" };
  return { caseId: base, role: "label" };
}

// --------------------------- Grouping ---------------------------

/** Group a flat file list into cases using nnU-Net naming + a split hint. */
export function groupFilesIntoCases(
  files: SourceFile[],
  meta: DatasetMeta
): { cases: DatasetCase[]; unassigned: SourceFile[] } {
  const byCase = new Map<string, DatasetCase>();
  const unassigned: SourceFile[] = [];

  const ensure = (caseId: string, split: Split): DatasetCase => {
    let c = byCase.get(caseId);
    if (!c) {
      c = { caseId, split, channelFiles: {}, labelFile: undefined };
      byCase.set(caseId, c);
    }
    return c;
  };

  for (const f of files) {
    const parsed =
      f.caseId != null && f.role !== "unknown"
        ? { caseId: f.caseId, channelIndex: f.channelIndex, role: f.role }
        : parseFilename(f.filename, meta.fileEnding);

    if (!parsed) {
      unassigned.push(f);
      continue;
    }
    const c = ensure(parsed.caseId, f.split);
    if (parsed.role === "image" && parsed.channelIndex != null) {
      c.channelFiles[parsed.channelIndex] = { ...f, role: "image", caseId: parsed.caseId, channelIndex: parsed.channelIndex };
    } else {
      c.labelFile = { ...f, role: "label", caseId: parsed.caseId };
    }
  }

  return { cases: [...byCase.values()], unassigned };
}

// ------------------- RedBrick Items List import -------------------

/** Minimal shape of a RedBrick Items List entry (standard format). */
export interface RedBrickEntry {
  name: string;
  series?: { items: string | string[]; name?: string }[];
  items?: string[]; // simplified auto-split format
  preAssign?: Record<string, string | string[]>;
}

/**
 * Map a RedBrick Items List into nnU-Net source files.
 * Each entry -> one case; each series -> one channel (in array order).
 * Mirrors the RedBrick spec: Task=Case, series=Channel.
 */
export function fromRedBrickItemsList(
  entries: RedBrickEntry[],
  ending: FileEnding,
  split: Split = "train"
): SourceFile[] {
  const out: SourceFile[] = [];
  entries.forEach((entry, ei) => {
    const seriesList =
      entry.series ??
      (entry.items ? entry.items.map((it) => ({ items: it })) : []);
    seriesList.forEach((s, ci) => {
      const first = Array.isArray(s.items) ? s.items[0] : s.items;
      out.push({
        id: `rb-${ei}-${ci}-${Math.random().toString(36).slice(2, 7)}`,
        filename: `${entry.name}_${pad4(ci)}${ending}`,
        sizeBytes: 0,
        role: "image",
        caseId: entry.name,
        channelIndex: ci,
        split,
      });
      void first;
    });
  });
  return out;
}

// --------------------------- Validation ---------------------------

export function validateDataset(
  meta: DatasetMeta,
  cases: DatasetCase[]
): ValidationReport {
  const issues: ValidationIssue[] = [];
  const err = (code: string, message: string, caseId?: string) =>
    issues.push({ level: "error", code, message, caseId });
  const warn = (code: string, message: string, caseId?: string) =>
    issues.push({ level: "warning", code, message, caseId });

  // Dataset id / name
  if (!Number.isInteger(meta.id) || meta.id < 1 || meta.id > 999)
    err("dataset.id", "Dataset ID must be an integer between 1 and 999.");
  if (meta.id >= 1 && meta.id <= 10)
    warn("dataset.idReserved", "IDs 001–010 are reserved for the Medical Segmentation Decathlon.");
  if (!meta.name.trim()) err("dataset.name", "Dataset name is required.");
  if (/\s/.test(meta.name)) warn("dataset.nameSpaces", "Dataset name should not contain spaces (use CamelCase).");

  // Channels
  if (meta.channels.length === 0) err("channels.empty", "Define at least one input channel.");
  const chIdx = meta.channels.map((c) => c.index).sort((a, b) => a - b);
  chIdx.forEach((v, i) => {
    if (v !== i) err("channels.consecutive", "Channel indices must be consecutive starting at 0.");
  });
  const chNames = meta.channels.map((c) => c.name.trim());
  if (chNames.some((n) => !n)) err("channels.name", "Every channel needs a name.");
  if (new Set(chNames).size !== chNames.length) warn("channels.dupName", "Channel names should be unique.");

  // Labels
  if (!meta.labels.some((l) => l.value === 0))
    err("labels.background", "Labels must include background with value 0.");
  const labelVals = meta.labels.map((l) => l.value).sort((a, b) => a - b);
  labelVals.forEach((v, i) => {
    if (v !== i) err("labels.consecutive", "Label values must be consecutive integers starting at 0.");
  });
  const labelNames = meta.labels.map((l) => l.name.trim());
  if (labelNames.some((n) => !n)) err("labels.name", "Every label needs a name.");
  if (new Set(labelNames).size !== labelNames.length) err("labels.dupName", "Label names must be unique.");

  // Dimensionality / format coherence
  if (!endingsForDim(meta.dim).includes(meta.fileEnding))
    warn("format.dimMismatch", `${meta.fileEnding} is unusual for a ${meta.dim.toUpperCase()} dataset.`);
  if (meta.rgb) {
    const names = meta.channels.map((c) => c.name.trim().toUpperCase()).join(",");
    if (meta.channels.length !== 3 || names !== "R,G,B")
      warn("rgb.channels", "RGB mode expects exactly three channels named R, G, B.");
    if (![".png", ".bmp", ".tif", ".tiff"].includes(meta.fileEnding))
      err("rgb.format", "RGB mode requires a 2D image format (.png, .bmp, .tif).");
  }

  // Cases
  const trainCases = cases.filter((c) => c.split === "train");
  if (trainCases.length === 0) err("cases.noTrain", "At least one training case is required.");

  const ids = cases.map((c) => c.caseId);
  const dup = ids.filter((id, i) => ids.indexOf(id) !== i);
  if (dup.length) err("cases.dupId", `Duplicate case identifiers: ${[...new Set(dup)].join(", ")}.`);

  const reqChannels = fileChannels(meta);
  for (const c of cases) {
    for (const ch of reqChannels) {
      if (!c.channelFiles[ch.index])
        err("case.missingChannel", meta.rgb ? "Missing the RGB image file." : `Missing channel ${pad4(ch.index)} (${ch.name}).`, c.caseId);
    }
    if (c.split === "train" && !c.labelFile)
      err("case.missingLabel", "Training case has no label file (labelsTr).", c.caseId);
    if (c.split === "test" && c.labelFile)
      warn("case.testLabel", "Test case has a label; labels are ignored for the test set (imagesTs).", c.caseId);

    // PNG inspection (only for files whose headers we actually read)
    const dims: { w: number; h: number; name: string }[] = [];
    for (const ch of reqChannels) {
      const f = c.channelFiles[ch.index];
      if (f?.png) dims.push({ w: f.png.width, h: f.png.height, name: f.filename });
    }
    if (c.labelFile?.png) {
      dims.push({ w: c.labelFile.png.width, h: c.labelFile.png.height, name: c.labelFile.filename });
      if (c.labelFile.png.colorType === 2 || c.labelFile.png.colorType === 6)
        warn("label.rgb", `Label ${c.labelFile.filename} is ${pngColorTypeName(c.labelFile.png.colorType)}; labels must be single-channel integer masks.`, c.caseId);
    }
    if (dims.length >= 2) {
      const { w, h } = dims[0];
      if (dims.some((d) => d.w !== w || d.h !== h))
        err("case.geometry", `Image and label dimensions differ within the case (${dims.map((d) => `${d.w}×${d.h}`).join(", ")}).`, c.caseId);
    }
  }

  // File format consistency + lossy guard
  const allFiles = cases.flatMap((c) => [
    ...Object.values(c.channelFiles).filter(Boolean) as SourceFile[],
    ...(c.labelFile ? [c.labelFile] : []),
  ]);
  for (const f of allFiles) {
    if (/\.jpe?g$/i.test(f.filename))
      err("file.lossy", `Lossy format not allowed: ${f.filename}. Use a lossless format.`);
    if (!f.filename.endsWith(meta.fileEnding))
      err("file.ending", `File ending must be ${meta.fileEnding}: ${f.filename}.`);
  }

  const errorCount = issues.filter((i) => i.level === "error").length;
  const warningCount = issues.filter((i) => i.level === "warning").length;
  return { issues, errorCount, warningCount, ok: errorCount === 0 };
}

// --------------------- Outputs: json + tree ---------------------

export function generateDatasetJson(
  meta: DatasetMeta,
  cases: DatasetCase[]
): Record<string, unknown> {
  const channel_names: Record<string, string> = {};
  [...meta.channels]
    .sort((a, b) => a.index - b.index)
    .forEach((c) => (channel_names[String(c.index)] = c.isCT ? "CT" : c.name));

  const labels: Record<string, number> = {};
  [...meta.labels]
    .sort((a, b) => a.value - b.value)
    .forEach((l) => (labels[l.name || `label_${l.value}`] = l.value));

  return {
    channel_names,
    labels,
    numTraining: cases.filter((c) => c.split === "train").length,
    file_ending: meta.fileEnding,
  };
}

/** Source -> target path plan (the rename/move manifest). */
export function planFolderTree(meta: DatasetMeta, cases: DatasetCase[]): PlannedFile[] {
  const root = datasetFolderName(meta);
  const out: PlannedFile[] = [];
  const chans = fileChannels(meta);
  for (const c of cases) {
    const imgDir = c.split === "train" ? "imagesTr" : "imagesTs";
    for (const ch of chans) {
      const f = c.channelFiles[ch.index];
      if (f)
        out.push({
          source: f,
          targetPath: `${root}/${imgDir}/${c.caseId}_${pad4(ch.index)}${meta.fileEnding}`,
        });
    }
    if (c.split === "train" && c.labelFile)
      out.push({
        source: c.labelFile,
        targetPath: `${root}/labelsTr/${c.caseId}${meta.fileEnding}`,
      });
  }
  return out;
}

// --------------------- Execution: script emitters ---------------------

export type Operation = "copy" | "move";
export type Target = "local" | "s3" | "gcs" | "azure";

export interface ExecOptions {
  operation: Operation;
  target: Target;
  /** Where the raw files currently live (filesystem path or cloud URI). */
  sourceRoot: string;
  /** nnUNet_raw destination (filesystem path or cloud URI). */
  destRoot: string;
}

const uniqueDirs = (plan: PlannedFile[]): string[] =>
  [...new Set(plan.map((p) => p.targetPath.split("/").slice(0, -1).join("/")))];

/** POSIX bash script: mkdir tree, write dataset.json, copy/move each file. */
export function generateBashScript(
  meta: DatasetMeta,
  plan: PlannedFile[],
  datasetJson: Record<string, unknown>,
  opts: ExecOptions
): string {
  const DS = datasetFolderName(meta);
  const op = opts.operation === "move" ? "mv" : "cp -p";
  const L: string[] = [
    "#!/usr/bin/env bash",
    "# Sắp xếp tệp thô theo cấu trúc thư mục chuẩn. Kiểm tra trước khi chạy.",
    "set -euo pipefail",
    "",
    `SOURCE_DIR="\${1:-${opts.sourceRoot || "./incoming"}}"`,
    `DEST_ROOT="\${nnUNet_raw:-${opts.destRoot || "./nnUNet_raw"}}"`,
    "",
    "# 1. folder structure",
  ];
  uniqueDirs(plan).forEach((d) => L.push(`mkdir -p "$DEST_ROOT/${d}"`));
  L.push("", "# 2. dataset.json", `cat > "$DEST_ROOT/${DS}/dataset.json" <<'JSON'`, JSON.stringify(datasetJson, null, 2), "JSON", "");
  L.push("# 3. transfer files (tab-separated src<TAB>dst)", `while IFS=$'\\t' read -r src dst; do`, `  [ -z "$src" ] && continue`, `  ${op} "$SOURCE_DIR/$src" "$DEST_ROOT/$dst"`, `done <<'MANIFEST'`);
  plan.forEach((p) => L.push(`${p.source.filename}\t${p.targetPath}`));
  L.push("MANIFEST", "", `echo "Organized ${plan.length} files into $DEST_ROOT/${DS}"`);
  return L.join("\n");
}

/** Cross-platform Python (shutil) script. */
export function generatePythonScript(
  meta: DatasetMeta,
  plan: PlannedFile[],
  datasetJson: Record<string, unknown>,
  opts: ExecOptions
): string {
  const DS = datasetFolderName(meta);
  const manifest = plan.map((p) => `    (${JSON.stringify(p.source.filename)}, ${JSON.stringify(p.targetPath)}),`).join("\n");
  return [
    "#!/usr/bin/env python3",
    '"""Sắp xếp tệp thô theo cấu trúc thư mục chuẩn. Kiểm tra trước khi chạy."""',
    "import os, sys, json, shutil",
    "",
    `SOURCE_DIR = sys.argv[1] if len(sys.argv) > 1 else ${JSON.stringify(opts.sourceRoot || "./incoming")}`,
    `DEST_ROOT = os.environ.get("nnUNet_raw", ${JSON.stringify(opts.destRoot || "./nnUNet_raw")})`,
    `OPERATION = ${JSON.stringify(opts.operation)}  # "copy" or "move"`,
    "",
    "DATASET_JSON = " + JSON.stringify(datasetJson, null, 2),
    "",
    "MANIFEST = [",
    manifest,
    "]",
    "",
    `root = os.path.join(DEST_ROOT, ${JSON.stringify(DS)})`,
    "for d in ('imagesTr', 'labelsTr', 'imagesTs'):",
    "    os.makedirs(os.path.join(root, d), exist_ok=True)",
    "with open(os.path.join(root, 'dataset.json'), 'w') as f:",
    "    json.dump(DATASET_JSON, f, indent=2)",
    "",
    "xfer = shutil.move if OPERATION == 'move' else shutil.copy2",
    "for src, dst in MANIFEST:",
    "    s = os.path.join(SOURCE_DIR, src)",
    "    t = os.path.join(DEST_ROOT, dst)",
    "    os.makedirs(os.path.dirname(t), exist_ok=True)",
    "    xfer(s, t)",
    "",
    "print(f'Organized {len(MANIFEST)} files into {root}')",
  ].join("\n");
}

/** Cloud CLI script (aws s3 / gsutil / azcopy). */
export function generateCloudScript(
  meta: DatasetMeta,
  plan: PlannedFile[],
  datasetJson: Record<string, unknown>,
  opts: ExecOptions
): string {
  const DS = datasetFolderName(meta);
  const p = opts.target;
  const xfer =
    p === "s3" ? (opts.operation === "move" ? "aws s3 mv" : "aws s3 cp")
    : p === "gcs" ? (opts.operation === "move" ? "gsutil mv" : "gsutil cp")
    : "azcopy copy";
  const up = p === "s3" ? "aws s3 cp" : p === "gcs" ? "gsutil cp" : "azcopy copy";
  const defSrc = p === "s3" ? "s3://my-bucket/incoming" : p === "gcs" ? "gs://my-bucket/incoming" : "https://acct.blob.core.windows.net/incoming";
  const defDst = p === "s3" ? "s3://my-bucket/nnUNet_raw" : p === "gcs" ? "gs://my-bucket/nnUNet_raw" : "https://acct.blob.core.windows.net/nnUNet_raw";
  const L: string[] = [
    "#!/usr/bin/env bash",
    `# Chuyển tệp thô theo cấu trúc thư mục chuẩn lên ${p.toUpperCase()}. Kiểm tra trước khi chạy.`,
    "set -euo pipefail",
    "",
    `SRC="${opts.sourceRoot || defSrc}"`,
    `DST="${opts.destRoot || defDst}"`,
    "",
    "# dataset.json",
    "cat > dataset.json <<'JSON'",
    JSON.stringify(datasetJson, null, 2),
    "JSON",
    `${up} "dataset.json" "$DST/${DS}/dataset.json"`,
    "",
    "# files",
  ];
  plan.forEach((pf) => L.push(`${xfer} "$SRC/${pf.source.filename}" "$DST/${pf.targetPath}"`));
  if (opts.operation === "move" && p === "azure")
    L.push("", "# NOTE: azcopy has no atomic move — delete sources separately if required.");
  L.push("", `echo "Transferred ${plan.length} files to $DST/${DS}"`);
  return L.join("\n");
}
