import { useState, useCallback, useRef, useEffect } from "react";
import vtkImageData from "@kitware/vtk.js/Common/DataModel/ImageData";
import vtkDataArray from "@kitware/vtk.js/Common/Core/DataArray";

interface ParsedSlice {
  data: Uint8Array | Int16Array;
  width: number;
  height: number;
  bitsAllocated: 8 | 16;
}

interface DicomData {
  imageData: any;
  slices: ParsedSlice[];
  width: number;
  height: number;
  depth: number;
}

// Simple DICOM parser for raw pixel data
const parseDicomBuffer = (buffer: Uint8Array): ParsedSlice | null => {
  try {
    const hasDicmPrefix =
      buffer.length > 132 &&
      buffer[128] === 68 &&
      buffer[129] === 73 &&
      buffer[130] === 67 &&
      buffer[131] === 77;

    let offset = hasDicmPrefix ? 132 : 0;
    let width = 512;
    let height = 512;
    let bitsAllocated: 8 | 16 = 16;
    let pixelDataOffset = -1;
    let pixelDataLength = 0;

    while (offset < buffer.length - 8) {
      const group = buffer[offset] | (buffer[offset + 1] << 8);
      const element = buffer[offset + 2] | (buffer[offset + 3] << 8);
      offset += 4;

      let length: number;
      const vr = String.fromCharCode(buffer[offset], buffer[offset + 1]);

      if (["OB", "OW", "OF", "SQ", "UC", "UN", "UR", "UT"].includes(vr)) {
        offset += 4;
        length =
          buffer[offset] |
          (buffer[offset + 1] << 8) |
          (buffer[offset + 2] << 16) |
          (buffer[offset + 3] << 24);
        offset += 4;
      } else if (vr.match(/[A-Z]{2}/)) {
        offset += 2;
        length = buffer[offset] | (buffer[offset + 1] << 8);
        offset += 2;
      } else {
        length =
          buffer[offset - 4] |
          (buffer[offset - 3] << 8) |
          (buffer[offset - 2] << 16) |
          (buffer[offset - 1] << 24);
      }

      if (group === 0x0028 && element === 0x0010) {
        height = buffer[offset] | (buffer[offset + 1] << 8);
      } else if (group === 0x0028 && element === 0x0011) {
        width = buffer[offset] | (buffer[offset + 1] << 8);
      } else if (group === 0x0028 && element === 0x0100) {
        const v = buffer[offset] | (buffer[offset + 1] << 8);
        bitsAllocated = v === 8 ? 8 : 16;
      } else if (group === 0x7fe0 && element === 0x0010) {
        pixelDataOffset = offset;
        pixelDataLength = length;
        break;
      }

      if (length === 0xffffffff) break;
      offset += length;
    }

    if (pixelDataOffset === -1) {
      const expectedSize = width * height * (bitsAllocated / 8);
      if (buffer.length >= expectedSize) {
        pixelDataOffset = buffer.length - expectedSize;
        pixelDataLength = expectedSize;
      } else {
        return null;
      }
    }

    if (bitsAllocated === 16) {
      // Interpret as SIGNED Int16 (typical CT storage), keep raw values for window/level
      const pixelBuffer = buffer.slice(pixelDataOffset, pixelDataOffset + pixelDataLength);
      const out = new Int16Array(width * height);
      for (let i = 0; i < width * height; i++) {
        const low = pixelBuffer[i * 2] ?? 0;
        const high = pixelBuffer[i * 2 + 1] ?? 0;
        let value = low | (high << 8);
        if (value & 0x8000) value = value - 0x10000; // signed
        out[i] = value;
      }
      return { data: out, width, height, bitsAllocated: 16 };
    }

    // 8-bit fallback
    const pixelBuffer = buffer.slice(
      pixelDataOffset,
      pixelDataOffset + Math.min(pixelDataLength, width * height)
    );
    const out = new Uint8Array(width * height);
    out.set(pixelBuffer.subarray(0, out.length));
    return { data: out, width, height, bitsAllocated: 8 };
  } catch (error) {
    console.error("Error parsing DICOM:", error);
    return null;
  }
};

export const useDicomData = (dicomFiles: File[]) => {
  const [dicomData, setDicomData] = useState<DicomData | null>(null);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const lastFilesRef = useRef<string>("");

  const loadDicomFiles = useCallback(async (files: File[]) => {
    if (files.length === 0) {
      setDicomData(null);
      return;
    }

    // Create a unique key from file names and sizes to detect changes
    const filesKey = files
      .map((f) => `${f.name}-${f.size}-${f.lastModified}`)
      .sort()
      .join("|");

    // Skip if we already loaded these files
    if (filesKey === lastFilesRef.current && dicomData) {
      return;
    }

    lastFilesRef.current = filesKey;
    setIsLoading(true);
    setError(null);

    try {
      const sortedFiles = [...files].sort((a, b) =>
        a.name.localeCompare(b.name, undefined, { numeric: true })
      );

      const slices: ParsedSlice[] = [];

      for (const file of sortedFiles) {
        const arrayBuffer = await file.arrayBuffer();
        const parsed = parseDicomBuffer(new Uint8Array(arrayBuffer));
        if (parsed) {
          slices.push(parsed);
        }
      }

      if (slices.length === 0) {
        setError("No valid DICOM data found");
        setIsLoading(false);
        return;
      }

      const width = slices[0].width;
      const height = slices[0].height;
      const depth = slices.length;

      const is16Bit = slices[0].bitsAllocated === 16;
      const voxelCount = width * height * depth;

      const volumeData = is16Bit ? new Int16Array(voxelCount) : new Uint8Array(voxelCount);

      for (let z = 0; z < depth; z++) {
        const slice = slices[z];
        const sliceOffset = z * width * height;
        if (is16Bit) {
          (volumeData as Int16Array).set(slice.data as Int16Array, sliceOffset);
        } else {
          (volumeData as Uint8Array).set(slice.data as Uint8Array, sliceOffset);
        }
      }

      const imageData = vtkImageData.newInstance();
      imageData.setDimensions(width, height, depth);
      imageData.setSpacing([1.0, 1.0, 1.0]);

      const scalars = vtkDataArray.newInstance({
        name: "Scalars",
        numberOfComponents: 1,
        values: volumeData as any,
      });

      imageData.getPointData().setScalars(scalars);

      setDicomData({
        imageData,
        slices,
        width,
        height,
        depth,
      });
    } catch (err) {
      console.error("Error loading DICOM files:", err);
      setError("Error loading DICOM files");
    } finally {
      setIsLoading(false);
    }
  }, [dicomData]);

  useEffect(() => {
    loadDicomFiles(dicomFiles);
  }, [dicomFiles, loadDicomFiles]);

  return { dicomData, isLoading, error };
};
