import { type Dispatch, type ReactNode } from 'react';
import { ChevronDown, ChevronUp, Plus, Trash2 } from 'lucide-react';
import {
  Badge,
  Button,
  Checkbox,
  Input,
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
  cn,
} from '@ump/shared';
import { datasetFolderName, pad4, type FileEnding } from '../../domain/nnunet/core';
import {
  endingsForState,
  type DatasetKind,
  type NnunetAction,
  type NnunetWizardState,
} from '../../domain/nnunet/wizardState';

type StepProps = { state: NnunetWizardState; dispatch: Dispatch<NnunetAction> };

// ---------- shared step primitives (UMP design system) ----------

export function StepHead({ title, sub }: { title: string; sub: string }) {
  return (
    <div className="mb-5">
      <h2 className="font-serif text-lg font-semibold text-foreground">{title}</h2>
      <p className="mt-0.5 max-w-2xl text-sm text-muted-foreground">{sub}</p>
    </div>
  );
}

export function Field({ label, hint, children }: { label: string; hint?: string; children: ReactNode }) {
  return (
    <label className="block">
      <div className="mb-1.5 text-xs font-medium text-foreground">{label}</div>
      {children}
      {hint && <div className="mt-1 text-[11px] text-muted-foreground">{hint}</div>}
    </label>
  );
}

// ---------- Step 0 — Dataset identity ----------

const KIND_LABELS: Record<DatasetKind, string> = {
  volume3d: 'Khối 3D — NIfTI / NRRD / MHA (mỗi kênh một tệp)',
  gray2d: 'Ảnh xám 2D — PNG / BMP / TIF (mỗi kênh một tệp)',
  rgb2d: 'Ảnh RGB 2D — PNG (ba kênh màu trong một tệp)',
};

export function DatasetStep({ state, dispatch }: StepProps) {
  const { meta } = state;
  const kind: DatasetKind = meta.rgb ? 'rgb2d' : meta.dim === '2d' ? 'gray2d' : 'volume3d';
  const config = meta.dim === '2d' ? '2d' : '3d_fullres';
  return (
    <div>
      <StepHead
        title="Định danh bộ dữ liệu"
        sub="Chọn loại dữ liệu rồi đặt tên. Bộ dữ liệu được lưu theo tên DatasetXXX_Tên."
      />
      <Field label="Loại dữ liệu" hint="Quyết định số chiều, định dạng tệp và cách kênh ánh xạ vào tệp.">
        <Select value={kind} onValueChange={(v) => dispatch({ type: 'SET_DATASET_KIND', kind: v as DatasetKind })}>
          <SelectTrigger className="max-w-md">
            <SelectValue />
          </SelectTrigger>
          <SelectContent>
            {(Object.keys(KIND_LABELS) as DatasetKind[]).map((k) => (
              <SelectItem key={k} value={k}>
                {KIND_LABELS[k]}
              </SelectItem>
            ))}
          </SelectContent>
        </Select>
      </Field>
      <div className="mt-4 grid max-w-xl gap-4 sm:grid-cols-2">
        <Field label="Mã bộ dữ liệu (ID)" hint="1–999. ID 001–010 dành cho Medical Segmentation Decathlon.">
          <Input
            type="number"
            min={1}
            max={999}
            value={meta.id}
            onChange={(e) => dispatch({ type: 'PATCH_META', patch: { id: parseInt(e.target.value || '0', 10) } })}
            className="font-mono"
          />
        </Field>
        <Field label="Tên bộ dữ liệu" hint="CamelCase, không khoảng trắng.">
          <Input
            value={meta.name}
            onChange={(e) => dispatch({ type: 'PATCH_META', patch: { name: e.target.value } })}
            placeholder="Prostate"
          />
        </Field>
        <Field label="Đuôi tệp" hint={meta.dim === '2d' ? 'Định dạng 2D không mất dữ liệu (không .jpg).' : 'Định dạng 3D không mất dữ liệu.'}>
          <Select value={meta.fileEnding} onValueChange={(v) => dispatch({ type: 'PATCH_META', patch: { fileEnding: v as FileEnding } })}>
            <SelectTrigger className="font-mono">
              <SelectValue />
            </SelectTrigger>
            <SelectContent>
              {endingsForState(state).map((e) => (
                <SelectItem key={e} value={e} className="font-mono">
                  {e}
                </SelectItem>
              ))}
            </SelectContent>
          </Select>
        </Field>
      </div>
      <div className="mt-6 flex flex-wrap items-center justify-between gap-2 rounded-lg border border-border bg-muted/40 p-3.5">
        <div>
          <div className="mb-1 text-[11px] uppercase tracking-wide text-muted-foreground">Thư mục kết quả</div>
          <div className="font-mono text-primary">{datasetFolderName(meta)}/</div>
        </div>
        <div className="text-right">
          <div className="mb-1 text-[11px] uppercase tracking-wide text-muted-foreground">Cấu hình huấn luyện</div>
          <div className="font-mono text-primary">{config}</div>
        </div>
      </div>
    </div>
  );
}

// ---------- Step 1 — Channels ----------

export function ChannelsStep({ state, dispatch }: StepProps) {
  const { meta } = state;
  if (meta.rgb) {
    return (
      <div>
        <StepHead title="Kênh đầu vào" sub="Chế độ RGB: ba kênh màu được lưu chung trong một ảnh cho mỗi case." />
        <div className="max-w-2xl space-y-2">
          {['R', 'G', 'B'].map((n, i) => (
            <div key={i} className="flex items-center gap-2 rounded-lg border border-border bg-card p-2 opacity-90">
              <Badge variant="outline" className="w-14 justify-center font-mono text-primary">{pad4(i)}</Badge>
              <Input value={n} disabled className="flex-1" />
              <span className="px-2 text-[11px] text-muted-foreground">đã khóa</span>
            </div>
          ))}
        </div>
        <p className="mt-4 max-w-2xl rounded-lg border border-border bg-muted/40 p-3 text-xs text-muted-foreground">
          Ngoại lệ RGB: <span className="font-mono">dataset.json</span> khai báo ba kênh{' '}
          <span className="font-mono">{'{0:R, 1:G, 2:B}'}</span> nhưng mỗi case chỉ là một tệp{' '}
          <span className="font-mono">{`{case}_0000${meta.fileEnding}`}</span>. Đổi loại dữ liệu sang{' '}
          <b>Ảnh xám 2D</b> để tự định nghĩa kênh.
        </p>
      </div>
    );
  }
  return (
    <div>
      <StepHead
        title="Kênh đầu vào"
        sub={
          meta.dim === '2d'
            ? 'Mỗi kênh là một ảnh cho mỗi case. Mọi case phải có cùng kênh, cùng thứ tự.'
            : 'Mỗi kênh là một ảnh đồng đăng ký cho mỗi case (vd T2, ADC). Thứ tự + số lượng phải giống nhau ở mọi case.'
        }
      />
      <div className="max-w-2xl space-y-2">
        {meta.channels.map((c, i) => (
          <div key={i} className="flex items-center gap-2 rounded-lg border border-border bg-card p-2">
            <Badge variant="outline" className="w-14 justify-center font-mono text-primary">{pad4(c.index)}</Badge>
            <Input
              value={c.name}
              onChange={(e) => dispatch({ type: 'UPDATE_CHANNEL', index: i, patch: { name: e.target.value } })}
              placeholder="Tên kênh (vd T2)"
              className="flex-1"
            />
            <label
              className="flex cursor-pointer select-none items-center gap-1.5 px-2 text-xs text-muted-foreground"
              title="CT dùng chuẩn hóa cường độ toàn cục; các kênh khác dùng z-score theo từng kênh."
            >
              <Checkbox checked={c.isCT} onCheckedChange={(v) => dispatch({ type: 'UPDATE_CHANNEL', index: i, patch: { isCT: v === true } })} /> CT
            </label>
            <div className="flex flex-col">
              <button type="button" onClick={() => dispatch({ type: 'MOVE_CHANNEL', index: i, dir: -1 })} className="text-muted-foreground hover:text-foreground">
                <ChevronUp className="h-3.5 w-3.5" />
              </button>
              <button type="button" onClick={() => dispatch({ type: 'MOVE_CHANNEL', index: i, dir: 1 })} className="text-muted-foreground hover:text-foreground">
                <ChevronDown className="h-3.5 w-3.5" />
              </button>
            </div>
            <button type="button" onClick={() => dispatch({ type: 'REMOVE_CHANNEL', index: i })} className="p-1 text-muted-foreground hover:text-destructive">
              <Trash2 className="h-4 w-4" />
            </button>
          </div>
        ))}
      </div>
      <Button variant="ghost" size="sm" className="mt-3 text-primary hover:bg-primary/10" onClick={() => dispatch({ type: 'ADD_CHANNEL' })}>
        <Plus className="mr-1.5 h-4 w-4" /> Thêm kênh
      </Button>
      <p className="mt-4 max-w-xl text-xs text-muted-foreground">
        Đánh dấu <b>CT</b> để chuẩn hóa toàn cục. Tương đương RedBrick: mỗi kênh là một <span className="font-mono">series</span> của Task.
      </p>
    </div>
  );
}

// ---------- Step 2 — Labels ----------

export function LabelsStep({ state, dispatch }: StepProps) {
  const { meta } = state;
  return (
    <div>
      <StepHead title="Nhãn phân vùng" sub="Các lớp số nguyên trong mặt nạ. Nền (background) là 0; các giá trị phải liên tiếp." />
      <div className="max-w-2xl space-y-2">
        {[...meta.labels]
          .sort((a, b) => a.value - b.value)
          .map((l) => {
            const i = meta.labels.indexOf(l);
            const isBg = l.value === 0;
            return (
              <div key={i} className="flex items-center gap-2 rounded-lg border border-border bg-card p-2">
                <Badge variant="outline" className="w-10 justify-center font-mono">{l.value}</Badge>
                <Input
                  value={l.name}
                  onChange={(e) => dispatch({ type: 'UPDATE_LABEL', index: i, patch: { name: e.target.value } })}
                  placeholder={isBg ? 'background' : 'tên lớp (vd tumor)'}
                  className={cn('flex-1', isBg && 'text-muted-foreground')}
                  disabled={isBg}
                />
                {isBg ? (
                  <span className="px-2 text-[11px] text-muted-foreground">bắt buộc</span>
                ) : (
                  <button type="button" onClick={() => dispatch({ type: 'REMOVE_LABEL', index: i })} className="p-1 text-muted-foreground hover:text-destructive">
                    <Trash2 className="h-4 w-4" />
                  </button>
                )}
              </div>
            );
          })}
      </div>
      <Button variant="ghost" size="sm" className="mt-3 text-primary hover:bg-primary/10" onClick={() => dispatch({ type: 'ADD_LABEL' })}>
        <Plus className="mr-1.5 h-4 w-4" /> Thêm nhãn
      </Button>
      <p className="mt-4 max-w-xl text-xs text-muted-foreground">
        Tương đương RedBrick: mỗi nhãn ↔ một lớp phân vùng trong Taxonomy; Ground Truth trở thành mặt nạ trong{' '}
        <span className="font-mono">labelsTr</span>.
      </p>
    </div>
  );
}
