import { useEffect, useMemo, useState } from 'react';
import { UnifiedQuadViewRenderer, loadNiftiImageData, labelValues, extractBinaryLabel } from '@ump/shared/viewer';
import type { Annotation, AnnotationTool, OrganMaskData } from '@ump/shared/viewer';
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogTitle,
  TooltipProvider,
  Button,
  fileDownloadUrl,
  fetchAsFile,
} from '@ump/shared';
import { MousePointer2, Square, Circle, PenLine, Brush, Hexagon, Trash2, Eye, EyeOff, Loader2, Pencil, Check, X } from 'lucide-react';
import { toast } from 'sonner';

import { organDisplayName } from './totalSegmentatorLabels';

/** A linked organ mask available to overlay on the parent image (Phase D). */
export interface AvailableMask {
  id: string;
  organLabel: string;
  logicalPath: string;
  /** When true, this is a single multi-label volume (e.g. TotalSegmentator): the viewer loads it
   *  once and expands it into one toggleable overlay per distinct integer label. */
  multiLabel?: boolean;
}

/** One toggleable overlay row — a whole single-file mask, or one value of a multi-label volume. */
interface OrganEntry {
  key: string;
  name: string;
  maskId: string;
  /** Set only for a multi-label sub-organ: the integer label value to extract. */
  value?: number;
}

interface DatasetFileViewerDialogProps {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  file: File | null;
  fileName: string;
  datasetId: string;
  /** Organ masks linked to this image; each can be toggled as a coloured 3D overlay. */
  masks?: AvailableMask[];
  /** Per-dataset integer-label → name map; overrides the built-in TotalSegmentator names. */
  labelMap?: Record<string, string>;
  /** Persist a renamed multi-label organ (value → name). Enables the inline rename control. */
  onRenameLabel?: (value: number, name: string) => void;
}

// Sensible soft-tissue defaults; the viewer detects DICOM vs NIfTI from the file name.
const DEFAULTS = { windowWidth: 400, windowLevel: 40, opacity: 0.8 };

// Distinct overlay colours (0-255 RGB), assigned to organs by list order.
const ORGAN_PALETTE: [number, number, number][] = [
  [239, 68, 68], [34, 197, 94], [59, 130, 246], [234, 179, 8],
  [168, 85, 247], [236, 72, 153], [20, 184, 166], [249, 115, 22],
];

const TOOLS: { id: AnnotationTool; label: string; Icon: typeof Square }[] = [
  { id: 'none', label: 'Chọn', Icon: MousePointer2 },
  { id: 'bbox', label: 'Khung', Icon: Square },
  { id: 'points', label: 'Điểm', Icon: Circle },
  { id: 'pen', label: 'Bút', Icon: PenLine },
  { id: 'brush', label: 'Cọ', Icon: Brush },
  { id: 'polygon', label: 'Đa giác', Icon: Hexagon },
];
const TOOL_LABEL: Record<string, string> = {
  bbox: 'Khung', points: 'Điểm', pen: 'Bút', brush: 'Cọ', polygon: 'Đa giác',
};
const VIEW_LABEL: Record<string, string> = { axial: 'Axial', coronal: 'Coronal', sagittal: 'Sagittal' };

/**
 * Lazy-loaded host for the ImageHub 3D viewer (statically imports `@ump/shared/viewer`
 * → @kitware/vtk.js, so it must stay behind a React.lazy boundary). Full-screen
 * layout: top control bar, then a body row of the quad-view + an annotation side panel.
 */
export default function DatasetFileViewerDialog({
  open,
  onOpenChange,
  file,
  fileName,
  datasetId,
  masks = [],
  labelMap,
  onRenameLabel,
}: DatasetFileViewerDialogProps) {
  const [windowWidth, setWindowWidth] = useState(DEFAULTS.windowWidth);
  const [windowLevel, setWindowLevel] = useState(DEFAULTS.windowLevel);
  const [opacity, setOpacity] = useState(DEFAULTS.opacity);
  const [tool, setTool] = useState<AnnotationTool>('none');
  const [annotations, setAnnotations] = useState<Annotation[]>([]);

  // Organ-mask overlays. A single-file mask is one overlay; a multi-label mask (multiLabel) is
  // loaded once and expanded into one overlay per distinct integer label. Loaded vtkImageData is
  // keyed by entry key, so toggling off→on is instant after the first load.
  const [selectedIds, setSelectedIds] = useState<Set<string>>(new Set());
  const [organMaskById, setOrganMaskById] = useState<Record<string, OrganMaskData>>({});
  const [loadingId, setLoadingId] = useState<string | null>(null);
  const [labelVols, setLabelVols] = useState<Record<string, { imageData: any; values: number[] }>>({});
  const [enumerating, setEnumerating] = useState(false);
  // Inline rename of a multi-label organ: { value being edited, draft text }.
  const [editing, setEditing] = useState<{ value: number; draft: string } | null>(null);

  const colorFor = (i: number): [number, number, number] => ORGAN_PALETTE[i % ORGAN_PALETTE.length];

  // Load each multi-label mask once to enumerate the organs (integer labels) it contains.
  useEffect(() => {
    const pending = masks.filter((m) => m.multiLabel && !labelVols[m.id]);
    if (pending.length === 0) return;
    let cancelled = false;
    (async () => {
      setEnumerating(true);
      for (const m of pending) {
        try {
          const { url } = await fileDownloadUrl(datasetId, m.id);
          const maskFile = await fetchAsFile(url, m.logicalPath);
          const imageData = await loadNiftiImageData(maskFile);
          const values = labelValues(imageData);
          if (!cancelled) setLabelVols((prev) => ({ ...prev, [m.id]: { imageData, values } }));
        } catch {
          if (!cancelled) toast.error('Không tải được mặt nạ đa nhãn');
        }
      }
      if (!cancelled) setEnumerating(false);
    })();
    return () => {
      cancelled = true;
    };
    // labelVols omitted: the !labelVols[m.id] guard already prevents re-loading a cached volume.
  }, [masks, datasetId]); // eslint-disable-line react-hooks/exhaustive-deps

  // Flatten masks → toggleable organ entries (a multi-label mask expands per value), sorted by name.
  const entries = useMemo<OrganEntry[]>(() => {
    const out: OrganEntry[] = [];
    for (const m of masks) {
      if (m.multiLabel) {
        const lv = labelVols[m.id];
        if (lv) {
          for (const v of lv.values) {
            out.push({ key: `${m.id}:${v}`, name: labelMap?.[String(v)] || organDisplayName(v), maskId: m.id, value: v });
          }
        }
      } else {
        out.push({ key: m.id, name: m.organLabel || m.logicalPath, maskId: m.id });
      }
    }
    out.sort((a, b) => a.name.localeCompare(b.name, 'vi'));
    return out;
  }, [masks, labelVols, labelMap]);

  // The overlays handed to the renderer = the loaded vtkImageData for selected entries.
  const organMasks = useMemo(
    () =>
      entries
        .filter((e) => selectedIds.has(e.key))
        .map((e) => organMaskById[e.key])
        .filter(Boolean) as OrganMaskData[],
    [entries, selectedIds, organMaskById],
  );

  const toggleOrgan = async (entry: OrganEntry, i: number) => {
    if (selectedIds.has(entry.key)) {
      setSelectedIds((prev) => {
        const next = new Set(prev);
        next.delete(entry.key);
        return next;
      });
      return;
    }
    if (!organMaskById[entry.key]) {
      setLoadingId(entry.key);
      try {
        let imageData: any;
        if (entry.value != null) {
          // Multi-label: extract this organ's binary submask from the cached label volume (sync).
          const lv = labelVols[entry.maskId];
          if (!lv) throw new Error('multi-label volume not loaded');
          imageData = extractBinaryLabel(lv.imageData, entry.value);
        } else {
          const m = masks.find((x) => x.id === entry.maskId);
          if (!m) throw new Error('mask not found');
          const { url } = await fileDownloadUrl(datasetId, m.id);
          const maskFile = await fetchAsFile(url, m.logicalPath);
          imageData = await loadNiftiImageData(maskFile);
        }
        setOrganMaskById((prev) => ({
          ...prev,
          [entry.key]: { id: entry.key, organName: entry.name, imageData, color: colorFor(i) },
        }));
      } catch {
        toast.error('Không tải được mặt nạ cơ quan');
        setLoadingId(null);
        return;
      }
      setLoadingId(null);
    }
    setSelectedIds((prev) => new Set(prev).add(entry.key));
  };

  // Persist an inline organ rename (value → name) via the parent, then close the editor.
  const commitRename = () => {
    if (!editing || !onRenameLabel) return;
    const name = editing.draft.trim();
    if (name) onRenameLabel(editing.value, name);
    setEditing(null);
  };

  return (
    <Dialog open={open} onOpenChange={onOpenChange}>
      <DialogContent className="inset-0 flex h-full w-full max-w-none translate-x-0 translate-y-0 flex-col gap-0 overflow-hidden rounded-none border-0 p-0">
        {/* Top control bar */}
        <div className="flex shrink-0 flex-wrap items-center justify-between gap-x-6 gap-y-2 border-b bg-background px-4 py-3 pr-12">
          <DialogTitle className="truncate font-mono text-sm">{fileName}</DialogTitle>
          <DialogDescription className="sr-only">
            Trình xem ảnh y khoa 3D (DICOM/NIfTI) cho tệp {fileName}
          </DialogDescription>

          <div className="flex flex-wrap items-center gap-x-5 gap-y-2 text-xs text-muted-foreground">
            <label className="flex items-center gap-2">
              <span>Độ rộng cửa sổ</span>
              <input type="range" min={1} max={4000} step={1} value={windowWidth}
                onChange={(e) => setWindowWidth(Number(e.target.value))} className="w-28 accent-primary" />
              <span className="w-12 tabular-nums text-foreground">{windowWidth}</span>
            </label>
            <label className="flex items-center gap-2">
              <span>Mức cửa sổ</span>
              <input type="range" min={-1000} max={3000} step={1} value={windowLevel}
                onChange={(e) => setWindowLevel(Number(e.target.value))} className="w-28 accent-primary" />
              <span className="w-12 tabular-nums text-foreground">{windowLevel}</span>
            </label>
            <label className="flex items-center gap-2">
              <span>Độ mờ 3D</span>
              <input type="range" min={0} max={1} step={0.05} value={opacity}
                onChange={(e) => setOpacity(Number(e.target.value))} className="w-24 accent-primary" />
              <span className="w-10 tabular-nums text-foreground">{opacity.toFixed(2)}</span>
            </label>
          </div>
        </div>

        {/* Body: viewer + annotation side panel */}
        <div className="flex min-h-0 flex-1">
          <TooltipProvider>
            <div className="relative min-w-0 flex-1 bg-black">
              {open && file && (
                <UnifiedQuadViewRenderer
                  files={[file]}
                  windowWidth={windowWidth}
                  windowLevel={windowLevel}
                  opacity={opacity}
                  organMasks={organMasks}
                  annotationTool={tool}
                  annotations={annotations}
                  onAnnotationsChange={setAnnotations}
                />
              )}
            </div>
          </TooltipProvider>

          <aside className="flex w-72 shrink-0 flex-col border-l bg-background">
            {masks.length > 0 && (
              <div className="border-b">
                <div className="px-4 py-3">
                  <h3 className="text-sm font-semibold text-foreground">Cơ quan ({entries.length})</h3>
                  <p className="mt-0.5 text-[11px] text-muted-foreground">
                    {enumerating ? 'Đang nạp danh sách cơ quan…' : 'Chọn cơ quan để hiển thị lớp phủ 3D.'}
                  </p>
                </div>
                <ul className="max-h-56 space-y-1 overflow-y-auto px-2 pb-2">
                  {entries.map((e, i) => {
                    const on = selectedIds.has(e.key);
                    const loading = loadingId === e.key;
                    const [r, g, b] = colorFor(i);
                    const canRename = !!onRenameLabel && e.value != null;
                    const isEditing = editing != null && e.value != null && editing.value === e.value;
                    if (isEditing) {
                      return (
                        <li key={e.key}>
                          <div className="flex items-center gap-1 rounded-md border border-primary px-2 py-1">
                            <span
                              className="h-3 w-3 shrink-0 rounded-sm border border-black/10"
                              style={{ background: `rgb(${r}, ${g}, ${b})` }}
                            />
                            <input
                              autoFocus
                              value={editing.draft}
                              onChange={(ev) => setEditing({ value: editing.value, draft: ev.target.value })}
                              onKeyDown={(ev) => {
                                if (ev.key === 'Enter') commitRename();
                                else if (ev.key === 'Escape') setEditing(null);
                              }}
                              placeholder={`Nhãn ${e.value}`}
                              className="min-w-0 flex-1 bg-transparent text-xs text-foreground outline-none"
                            />
                            <button type="button" onClick={commitRename} title="Lưu"
                              className="shrink-0 text-primary hover:opacity-80">
                              <Check className="h-3.5 w-3.5" />
                            </button>
                            <button type="button" onClick={() => setEditing(null)} title="Hủy"
                              className="shrink-0 text-muted-foreground hover:text-destructive">
                              <X className="h-3.5 w-3.5" />
                            </button>
                          </div>
                        </li>
                      );
                    }
                    return (
                      <li key={e.key} className="flex items-center gap-1">
                        <button
                          type="button"
                          onClick={() => toggleOrgan(e, i)}
                          disabled={loading}
                          title={e.name}
                          className={`flex min-w-0 flex-1 items-center gap-2 rounded-md border px-2 py-1.5 text-left text-xs transition ${
                            on ? 'border-primary bg-primary/10' : 'border-border hover:bg-accent'
                          }`}
                        >
                          <span
                            className="h-3 w-3 shrink-0 rounded-sm border border-black/10"
                            style={{ background: `rgb(${r}, ${g}, ${b})`, opacity: on ? 1 : 0.4 }}
                          />
                          <span className="min-w-0 flex-1 truncate text-foreground">{e.name}</span>
                          {loading ? (
                            <Loader2 className="h-3.5 w-3.5 shrink-0 animate-spin text-muted-foreground" />
                          ) : on ? (
                            <Eye className="h-3.5 w-3.5 shrink-0 text-primary" />
                          ) : (
                            <EyeOff className="h-3.5 w-3.5 shrink-0 text-muted-foreground" />
                          )}
                        </button>
                        {canRename && (
                          <button
                            type="button"
                            onClick={() => setEditing({ value: e.value as number, draft: e.name })}
                            title="Đổi tên nhãn"
                            className="shrink-0 rounded-md border border-border p-1.5 text-muted-foreground transition hover:bg-accent hover:text-foreground"
                          >
                            <Pencil className="h-3.5 w-3.5" />
                          </button>
                        )}
                      </li>
                    );
                  })}
                </ul>
              </div>
            )}
            <div className="border-b px-4 py-3">
              <h3 className="text-sm font-semibold text-foreground">Công cụ chú thích</h3>
              <p className="mt-0.5 text-[11px] text-muted-foreground">
                Vẽ trên các lát cắt 2D (axial / coronal / sagittal).
              </p>
            </div>

            <div className="grid grid-cols-3 gap-2 p-3">
              {TOOLS.map(({ id, label, Icon }) => (
                <button
                  key={id}
                  type="button"
                  onClick={() => setTool(id)}
                  className={`flex flex-col items-center gap-1 rounded-md border px-2 py-2 text-[11px] transition ${
                    tool === id
                      ? 'border-primary bg-primary/10 text-primary'
                      : 'border-border text-muted-foreground hover:bg-accent'
                  }`}
                >
                  <Icon className="h-4 w-4" />
                  {label}
                </button>
              ))}
            </div>

            <div className="flex items-center justify-between border-y px-4 py-2">
              <span className="text-xs font-medium text-foreground">Chú thích ({annotations.length})</span>
              {annotations.length > 0 && (
                <Button
                  variant="ghost"
                  size="sm"
                  className="h-7 px-2 text-xs"
                  onClick={() => setAnnotations([])}
                >
                  <Trash2 className="mr-1 h-3.5 w-3.5" /> Xóa hết
                </Button>
              )}
            </div>

            <div className="min-h-0 flex-1 overflow-y-auto p-2">
              {annotations.length === 0 ? (
                <p className="px-2 py-6 text-center text-[11px] leading-relaxed text-muted-foreground">
                  Chọn một công cụ rồi vẽ trên lát cắt.
                  <br />
                  (Đa giác: bấm từng đỉnh, nhấn đúp để đóng.)
                </p>
              ) : (
                <ul className="space-y-1">
                  {annotations.map((a) => (
                    <li
                      key={a.id}
                      className="flex items-center justify-between gap-2 rounded border px-2 py-1 text-[11px]"
                    >
                      <span className="flex min-w-0 items-center gap-1.5">
                        <span className="h-2.5 w-2.5 shrink-0 rounded-full" style={{ background: a.color }} />
                        <span className="truncate">
                          {TOOL_LABEL[a.tool]} · {VIEW_LABEL[a.view]} · lát {a.sliceIndex + 1}
                        </span>
                      </span>
                      <button
                        type="button"
                        onClick={() => setAnnotations(annotations.filter((x) => x.id !== a.id))}
                        className="shrink-0 text-muted-foreground hover:text-destructive"
                        title="Xóa"
                      >
                        <Trash2 className="h-3.5 w-3.5" />
                      </button>
                    </li>
                  ))}
                </ul>
              )}
            </div>
          </aside>
        </div>

        {/* Footer hint */}
        <div className="shrink-0 border-t bg-background px-4 py-1.5 text-center text-[11px] text-muted-foreground">
          “Chọn”: cuộn chuyển lát · Ctrl/⌘ + cuộn: thu phóng · Nhấn đúp khung: phóng to · Công cụ vẽ: kéo trên lát cắt
        </div>
      </DialogContent>
    </Dialog>
  );
}
