import { useCallback, useMemo, useReducer } from 'react';
import { useNavigate, useParams } from 'react-router-dom';
import { useQuery } from '@tanstack/react-query';
import {
  AlertTriangle,
  ArrowLeft,
  ArrowRight,
  CheckCircle2,
  Database,
  FileCheck2,
  FolderTree,
  Layers,
  Tag,
  XCircle,
} from 'lucide-react';
import { Button, cn, getDataset } from '@ump/shared';
import {
  datasetFolderName,
  generateDatasetJson,
  groupFilesIntoCases,
  pad4,
  parsePngHeader,
  planFolderTree,
  validateDataset,
  type SourceFile,
  type Split,
} from '../../domain/nnunet/core';
import { initialNnunetState, NNUNET_STEPS, nnunetReducer } from '../../domain/nnunet/wizardState';
import { nnunetSourceFilesFromItemsList } from '../../domain/nnunet/fromItemsList';
import type { DataType } from '../../domain/dataTypes';
import { ChannelsStep, DatasetStep, LabelsStep } from './steps';
import { CasesStep, ExportStep } from './dataSteps';
import { OutputPane } from './OutputPane';

/** Synthesize a small demo dataset for the "Load example" button (channels × a few cases + a test case). */
function exampleFiles(meta: { fileEnding: string; rgb: boolean; dim: '2d' | '3d'; channels: { index: number }[] }): SourceFile[] {
  const e = meta.fileEnding;
  const out: SourceFile[] = [];
  const mk = (filename: string, split: Split): SourceFile => ({
    id: Math.random().toString(36).slice(2),
    filename,
    sizeBytes: (1024 * 1024 * (0.2 + Math.random() * 8)) | 0,
    role: 'unknown',
    split,
  });
  const chans = meta.rgb ? [0] : meta.channels.map((c) => c.index);
  const prefix = meta.rgb ? 'img' : meta.dim === '2d' ? 'scan' : 'case';
  ['00', '01', '02'].forEach((n) => {
    chans.forEach((ci) => out.push(mk(`${prefix}_${n}_${pad4(ci)}${e}`, 'train')));
    out.push(mk(`${prefix}_${n}${e}`, 'train'));
  });
  chans.forEach((ci) => out.push(mk(`${prefix}_50_${pad4(ci)}${e}`, 'test')));
  return out;
}

const STEP_META = [
  { label: 'Bộ dữ liệu', icon: Database },
  { label: 'Kênh', icon: Layers },
  { label: 'Nhãn', icon: Tag },
  { label: 'Các case', icon: FolderTree },
  { label: 'Kiểm tra & xuất', icon: FileCheck2 },
] as const;

/** nnU-Net v2 dataset-organizer wizard for a dataset — the new upload mode (project-workflow §data-import). */
export function NnunetOrganizerPage() {
  const { id = '' } = useParams();
  const navigate = useNavigate();
  const [state, dispatch] = useReducer(nnunetReducer, initialNnunetState);
  const { meta, files, step } = state;

  const dsQ = useQuery({ queryKey: ['imagehub', 'dataset', id], queryFn: () => getDataset(id), enabled: !!id });

  // Pure derivations (core.ts) — recomputed from {meta, files}; never stored in reducer state.
  const { cases, unassigned } = useMemo(() => groupFilesIntoCases(files, meta), [files, meta]);
  const report = useMemo(() => validateDataset(meta, cases), [meta, cases]);
  const datasetJson = useMemo(() => generateDatasetJson(meta, cases), [meta, cases]);
  const plan = useMemo(() => planFolderTree(meta, cases), [meta, cases]);
  const trainCount = cases.filter((c) => c.split === 'train').length;
  const testCount = cases.filter((c) => c.split === 'test').length;

  // Browser-coupled handlers (File objects, clipboard) the pure reducer can't own.
  const addFiles = useCallback((fileList: FileList | File[]) => {
    const next: SourceFile[] = Array.from(fileList).map((f) => ({
      id: Math.random().toString(36).slice(2),
      filename: f.name,
      sizeBytes: f.size ?? 0,
      role: 'unknown',
      split: 'train' as Split,
      blob: f,
    }));
    dispatch({ type: 'ADD_FILES', files: next });
    // Sniff PNG geometry so validateDataset can check dimensions/colour-type.
    next.forEach((sf) => {
      if (/\.png$/i.test(sf.filename) && sf.blob) {
        sf.blob
          .slice(0, 33)
          .arrayBuffer()
          .then((buf) => {
            const info = parsePngHeader(new Uint8Array(buf));
            if (info) dispatch({ type: 'SET_FILE_PNG', id: sf.id, png: info });
          })
          .catch(() => {});
      }
    });
  }, []);

  const importItemsList = useCallback(
    (json: string): string[] => {
      const dataType: DataType = meta.dim === '3d' ? 'nifti' : 'image2d';
      const r = nnunetSourceFilesFromItemsList(json, dataType, meta.fileEnding);
      if (r.ok) {
        dispatch({ type: 'ADD_FILES', files: r.files });
        return [];
      }
      return r.errors.map((e) => e.message);
    },
    [meta.dim, meta.fileEnding],
  );

  const loadExample = useCallback(() => {
    dispatch({ type: 'CLEAR_FILES' });
    dispatch({ type: 'ADD_FILES', files: exampleFiles(meta) });
  }, [meta]);

  const folderName = datasetFolderName(meta);
  const isLast = step === NNUNET_STEPS.length - 1;

  return (
    <div className="mx-auto max-w-[1400px] space-y-5">
      <Button variant="ghost" size="sm" onClick={() => navigate(`/dashboard/datasets/${id}`)}>
        <ArrowLeft className="mr-2 h-4 w-4" /> {dsQ.data?.name ?? 'Bộ dữ liệu'}
      </Button>

      {/* Header: title + resolved folder + validation status */}
      <div className="flex flex-wrap items-center justify-between gap-3">
        <div>
          <h1 className="font-serif text-2xl font-semibold text-foreground">Tổ chức dữ liệu</h1>
          <p className="text-sm text-muted-foreground">Sắp xếp dữ liệu hình ảnh y khoa theo cấu trúc thư mục chuẩn.</p>
        </div>
        <div className="flex items-center gap-4">
          <span className="rounded-md bg-muted px-2.5 py-1 font-mono text-sm text-foreground">{folderName}</span>
          <span className="flex items-center gap-1.5 text-sm font-medium">
            {report.errorCount ? (
              <>
                <XCircle className="h-4 w-4 text-destructive" />
                <span className="text-destructive">{report.errorCount} lỗi</span>
              </>
            ) : report.warningCount ? (
              <>
                <AlertTriangle className="h-4 w-4 text-amber-500" />
                <span className="text-amber-600">{report.warningCount} cảnh báo</span>
              </>
            ) : (
              <>
                <CheckCircle2 className="h-4 w-4 text-emerald-600" />
                <span className="text-emerald-600">Hợp lệ</span>
              </>
            )}
          </span>
        </div>
      </div>

      <div className="grid grid-cols-1 gap-5 lg:grid-cols-[180px_1fr_minmax(280px,360px)]">
        {/* Step nav */}
        <nav className="flex gap-1 overflow-x-auto lg:flex-col">
          {STEP_META.map((s, i) => {
            const Icon = s.icon;
            const active = i === step;
            const done = i < step;
            return (
              <button
                key={s.label}
                type="button"
                onClick={() => dispatch({ type: 'SET_STEP', step: i })}
                className={cn(
                  'flex items-center gap-2.5 whitespace-nowrap rounded-lg px-3 py-2.5 text-left text-sm transition',
                  active ? 'bg-primary text-primary-foreground shadow-sm' : 'text-muted-foreground hover:bg-muted hover:text-foreground',
                )}
              >
                <span
                  className={cn(
                    'grid h-5 w-5 shrink-0 place-items-center rounded-full font-mono text-[11px]',
                    active ? 'bg-primary-foreground/20' : done ? 'bg-emerald-100 text-emerald-700' : 'bg-muted text-muted-foreground',
                  )}
                >
                  {done ? <CheckCircle2 className="h-3.5 w-3.5" /> : i + 1}
                </span>
                <Icon className="h-4 w-4 shrink-0" />
                <span className="font-medium">{s.label}</span>
              </button>
            );
          })}
        </nav>

        {/* Step body + footer */}
        <main className="flex min-h-[460px] flex-col rounded-xl border border-border bg-card p-5 shadow-sm">
          <div className="flex-1">
            {step === 0 && <DatasetStep state={state} dispatch={dispatch} />}
            {step === 1 && <ChannelsStep state={state} dispatch={dispatch} />}
            {step === 2 && <LabelsStep state={state} dispatch={dispatch} />}
            {step === 3 && (
              <CasesStep
                datasetId={id}
                state={state}
                dispatch={dispatch}
                cases={cases}
                unassigned={unassigned}
                report={report}
                onAddFiles={addFiles}
                onImportItemsList={importItemsList}
                onLoadExample={loadExample}
              />
            )}
            {step === 4 && <ExportStep meta={meta} report={report} datasetJson={datasetJson} plan={plan} folderName={folderName} />}
          </div>
          <div className="mt-4 flex items-center justify-between border-t border-border pt-4">
            <Button variant="ghost" size="sm" disabled={step === 0} onClick={() => dispatch({ type: 'BACK' })}>
              <ArrowLeft className="mr-1.5 h-4 w-4" /> Quay lại
            </Button>
            {!isLast ? (
              <Button size="sm" onClick={() => dispatch({ type: 'NEXT' })}>
                Tiếp <ArrowRight className="ml-1.5 h-4 w-4" />
              </Button>
            ) : (
              <span
                className={cn(
                  'flex items-center gap-1.5 rounded-lg px-4 py-2 text-sm font-medium',
                  report.ok ? 'bg-emerald-50 text-emerald-700' : 'bg-amber-50 text-amber-700',
                )}
              >
                {report.ok ? (
                  <>
                    <CheckCircle2 className="h-4 w-4" /> Sẵn sàng xuất
                  </>
                ) : (
                  <>
                    <AlertTriangle className="h-4 w-4" /> Sửa lỗi để xuất
                  </>
                )}
              </span>
            )}
          </div>
        </main>

        <OutputPane
          folderName={folderName}
          plan={plan}
          datasetJson={datasetJson}
          trainCount={trainCount}
          testCount={testCount}
          channels={meta.channels.length}
          labels={meta.labels.length}
          dim={meta.dim}
          rgb={meta.rgb}
        />
      </div>

      {unassigned.length > 0 && step !== 3 && (
        <p className="text-xs text-amber-600">
          {unassigned.length} tệp chưa gom được vào case — xem bước “Các case”.
        </p>
      )}
    </div>
  );
}

export default NnunetOrganizerPage;
