import { useEffect, useState, useRef } from "react";
import marioImage from "./assets/mario.png";
import racoonImage from "./assets/racoon.png";
import cameramanImage from "./assets/cameraman.png";
import mandrillImage from "./assets/mandrill.png";
import motocrossImage from "./assets/motocross.png";
import peppersImage from "./assets/peppers.png";

import useConvolutionProcessing from "./useConvolutionProcessing.ts";
import type { ConvolutionRange } from "./useConvolutionProcessing.ts";
import { Card, Dropdown, Radio } from "@elvis/ui";

import { DEFAULT_KERNEL, KERNEL_PRESETS } from "./kernels.ts";
import type { Kernel } from "./kernels.ts";

const DEFAULT_IMAGE: string = mandrillImage;
const MIN_KERNEL_SIZE = 1;
const MAX_KERNEL_SIZE = 20;
const CUSTOM_2D_OPTION = "Custom2D...";
const CUSTOM_3D_OPTION = "Custom3D...";

export default function ConvolutionVisualizer() {
  const [rawInputImage, setRawInputImage] = useState<string>(DEFAULT_IMAGE);
  const [useColor, setUseColor] = useState<boolean>(true);
  const [kernel, setKernel] = useState<Kernel>(DEFAULT_KERNEL);
  const [stride, setStride] = useState(1);
  const [dilation, setDilation] = useState(1);
  const [padding, setPadding] = useState(0);

  const [inputImage, outputImage, outputRange] = useConvolutionProcessing(
    rawInputImage,
    kernel,
    useColor,
    stride,
    dilation,
    padding,
  );

  return (
    <div className="grid grid-cols-[minmax(0,7fr)_minmax(0,3fr)] gap-8 h-full min-h-0">
      <InputOutputViewer 
        input={inputImage}
        output={outputImage}
        outputRange={outputRange}
      />

      <Card className="flex flex-col gap-6 h-full min-h-0 overflow-auto">
        <ImageControls
          setImage={setRawInputImage}
          useColor={useColor}
          setUseColor={setUseColor}
        />
        <KernelEditor
          kernel={kernel}
          setKernel={setKernel}
          stride={stride}
          setStride={setStride}
          dilation={dilation}
          setDilation={setDilation}
          padding={padding}
          setPadding={setPadding}
        />
      </Card>
    </div>
  );
}

type InputOutputViewerProps = {
  input: string | null;
  output: string | null;
  outputRange: ConvolutionRange | null;
}
function InputOutputViewer({ input, output, outputRange }: InputOutputViewerProps) {
  const outputImageRef = useRef<HTMLImageElement | null>(null);
  const [outputDisplayWidth, setOutputDisplayWidth] = useState<number | null>(null);

  useEffect(() => {
    const image = outputImageRef.current;
    if (!image) return;

    function updateColorBarWidth() {
      if (!image || image.naturalWidth === 0 || image.naturalHeight === 0) return;
      const scale = Math.min(
        image.clientWidth / image.naturalWidth,
        image.clientHeight / image.naturalHeight,
      );
      setOutputDisplayWidth(image.naturalWidth * scale);
    }

    updateColorBarWidth();
    const observer = new ResizeObserver(updateColorBarWidth);
    observer.observe(image);
    return () => observer.disconnect();
  }, [output]);

  function formatScaleValue(value: number): string {
    if (Number.isInteger(value)) return value.toString();
    return Number(value.toPrecision(4)).toString();
  }

  return (
    <Card className="grid grid-cols-2 gap-4 h-full min-h-0">

      <div className="flex flex-col items-center justify-center min-h-0">
        <h2 className="text-lg font-bold mb-2">Input</h2>
        {input ? (
          <img src={input} alt="Input" className="flex-1 w-full h-full min-w-0 min-h-0 object-contain" />
        ) : (
          <p>Loading...</p>
        )}
        <div className="h-8 shrink-0" aria-hidden="true" />
      </div>

      <div className="flex flex-col items-center justify-center min-h-0">
        <h2 className="text-lg font-bold mb-2">Output</h2>
        {output ? (
          <>
            <div className="flex flex-1 items-center justify-center w-full min-w-0 min-h-0">
              <img
                ref={outputImageRef}
                src={output}
                alt="Output"
                className="flex-1 w-full h-full min-w-0 min-h-0 object-contain"
                onLoad={() => {
                  const image = outputImageRef.current;
                  if (!image || image.naturalWidth === 0 || image.naturalHeight === 0) return;
                  const scale = Math.min(
                    image.clientWidth / image.naturalWidth,
                    image.clientHeight / image.naturalHeight,
                  );
                  setOutputDisplayWidth(image.naturalWidth * scale);
                }}
              />
            </div>
            <div
              className="flex h-8 shrink-0 flex-col justify-end"
              style={outputDisplayWidth === null ? undefined : { width: outputDisplayWidth }}
              aria-label={outputRange
                ? `Grayscale color scale from ${outputRange.min} to ${outputRange.max}`
                : "Grayscale color scale"}
            >
              <div className="h-3 w-full border border-gray-400 bg-gradient-to-r from-black to-white" />
              <div className="flex justify-between text-xs leading-4 text-gray-700">
                <span>{outputRange ? formatScaleValue(outputRange.min) : ""}</span>
                <span>{outputRange ? formatScaleValue(outputRange.max) : ""}</span>
              </div>
            </div>
          </>
        ) : (
          <p>Processing...</p>
        )}
      </div>

    </Card>
  );
}

type ImageControlsProps = {
  setImage: (imageUrl: string) => void;
  useColor: boolean;
  setUseColor: (useColor: boolean) => void;
}
function ImageControls({ setImage, useColor, setUseColor }: ImageControlsProps) {
  const imageFileInputRef = useRef<HTMLInputElement | null>(null);
  const imageOptions = ["Upload...", "Enter URL...", "Mandrill", "Cameraman", "Motocross", "Peppers", "Super Mario", "Racoon"] as const;
  const [selectedImage, setSelectedImage] = useState<(typeof imageOptions)[number]>("Mandrill");
  const [imageUrl, setImageUrl] = useState<string>("");

  function loadImageUrl() {
    const trimmedUrl = imageUrl.trim();
    if (trimmedUrl) setImage(trimmedUrl);
  }

  function handleImageChange(option: (typeof imageOptions)[number]) {
    if (option === "Upload...") {
      imageFileInputRef.current?.click();
      return;
    }

    if (option === "Enter URL...") {
      setSelectedImage(option);
      return;
    }

    setSelectedImage(option);
    const presetImages = {
      Mandrill: mandrillImage,
			Cameraman: cameramanImage,
			Motocross: motocrossImage,
			Peppers: peppersImage,
      "Super Mario": marioImage,
      Racoon: racoonImage,
    };
    setImage(presetImages[option]);
  }

  return (
    <div className="flex flex-col gap-3">
      <input
        ref={imageFileInputRef}
        type="file"
        accept="image/*"
        onChange={(e) => {
          const file = e.target.files?.[0];
          if (!file) return;

          const reader = new FileReader();
          reader.onload = () => {
            const result = reader.result;
            if (typeof result !== "string") {
              return;
            }
            setImage(result);
          };
          reader.readAsDataURL(file);
          e.target.value = "";
        }}
        style={{ display: "none" }}
      />
      <div className="grid grid-cols-[4.5rem_minmax(0,1fr)] items-center gap-2">
        <span className="text-sm font-bold text-gray-700 text-right">Image</span>
        <div className="justify-self-start [&_label]:hidden [&_select]:px-2 [&_select]:py-1 [&_select]:text-xs">
          <Dropdown
            label=""
            options={imageOptions}
            activeOption={selectedImage}
            onChange={handleImageChange}
          />
        </div>
      </div>
      {selectedImage === "Enter URL..." && (
        <div className="grid grid-cols-[4.5rem_minmax(0,1fr)] items-center gap-2">
          <span aria-hidden="true" />
          <input
            type="url"
            value={imageUrl}
            placeholder="Enter image URL and press Enter"
            aria-label="Image URL"
            className="w-full rounded border border-gray-300 px-2 py-1 text-xs"
            onChange={(event) => setImageUrl(event.target.value)}
            onBlur={loadImageUrl}
            onKeyDown={(event) => {
              if (event.key === "Enter") loadImageUrl();
            }}
          />
        </div>
      )}
      <div className="grid grid-cols-[4.5rem_minmax(0,1fr)] items-center gap-2">
        <span className="text-sm font-bold text-gray-700 text-right">Color</span>
        <div className="justify-self-start text-sm">
          <Radio
            label=""
            options={["Grayscale", "Color"] as const}
            activeOption={useColor ? "Color" : "Grayscale"}
            onChange={(option) => setUseColor(option === "Color")}
          />
        </div>
      </div>
    </div>
  );
}

type KernelEditorProps = {
  kernel: Kernel;
  setKernel: (kernel: Kernel) => void;
  stride: number;
  setStride: (value: number) => void;
  dilation: number;
  setDilation: (value: number) => void;
  padding: number;
  setPadding: (value: number) => void;
}
function KernelEditor({
  kernel,
  setKernel,
  stride,
  setStride,
  dilation,
  setDilation,
  padding,
  setPadding,
}: KernelEditorProps) {
  const presetEntries = Object.entries(KERNEL_PRESETS);
  const presetGroups = [
    {
      label: "2D kernels",
      options: [
        CUSTOM_2D_OPTION,
        ...presetEntries
          .filter(([, preset]) => !Array.isArray(preset[0][0]))
          .map(([name]) => name),
      ],
    },
    {
      label: "3D kernels",
      options: [
        CUSTOM_3D_OPTION,
        ...presetEntries
          .filter(([, preset]) => Array.isArray(preset[0][0]))
          .map(([name]) => name),
      ],
    },
  ];
  const [selectedChannel, setSelectedChannel] = useState<number>(0);
  const [selectedPreset, setSelectedPreset] = useState<string>("Laplacian2D");

  function colorKernelToRepr(nextKernel: number[][][]): string[][][] {
    return nextKernel.map((channel) =>
      channel.map((row) => row.map((value) => value.toString())),
    );
  }

  function grayKernelToRepr(nextKernel: number[][]): string[][] {
    return nextKernel.map((row) => row.map((value) => value.toString()));
  }

  const [draftColorKernel, setDraftColorKernel] = useState<string[][][]>(
    colorKernelToRepr(KERNEL_PRESETS.Laplacian3D),
  );
  const [draftGrayKernel, setDraftGrayKernel] = useState<string[][]>(
    grayKernelToRepr(KERNEL_PRESETS.Laplacian2D),
  );

  useEffect(() => {
    if (Array.isArray(kernel[0][0])) {
      setDraftColorKernel(colorKernelToRepr(kernel as number[][][]));
    } else {
      setDraftGrayKernel(grayKernelToRepr(kernel as number[][]));
    }
  }, [kernel]);

  function parseCell(cell: string): number | null {
    const trimmed = cell.trim();
    if (trimmed === "") return null;
    const parsed = Number(trimmed);
    return Number.isFinite(parsed) ? parsed : null;
  }

  function cloneColorKernel(nextKernel: number[][][]): number[][][] {
    return nextKernel.map((channel) => channel.map((row) => [...row]));
  }

  function cloneGrayKernel(nextKernel: number[][]): number[][] {
    return nextKernel.map((row) => [...row]);
  }

  function handlePresetChange(option: string) {
    setSelectedPreset(option);
    const sourceKernel = option === CUSTOM_2D_OPTION
      ? KERNEL_PRESETS.Laplacian2D
      : option === CUSTOM_3D_OPTION
        ? KERNEL_PRESETS.Laplacian3D
        : KERNEL_PRESETS[option as keyof typeof KERNEL_PRESETS];

    if (Array.isArray(sourceKernel[0][0])) {
      const nextKernel = cloneColorKernel(sourceKernel as number[][][]);
      setDraftColorKernel(colorKernelToRepr(nextKernel));
      setKernel(nextKernel);
      setSelectedChannel(0);
    } else {
      const nextKernel = cloneGrayKernel(sourceKernel as number[][]);
      setDraftGrayKernel(grayKernelToRepr(nextKernel));
      setKernel(nextKernel);
    }
  }

  function handleKernelChange(channel: number, row: number, col: number, value: string) {
    if (Array.isArray(kernel[0][0])) {
      setSelectedPreset(CUSTOM_3D_OPTION);
      const nextDraft = draftColorKernel.map((c) => c.map((r) => [...r]));
      nextDraft[channel][row][col] = value;
      setDraftColorKernel(nextDraft);

      const parsedKernel = nextDraft.map((c) => c.map((r) => r.map(parseCell)));
      const isValid = parsedKernel.every((c) => c.every((r) => r.every((v) => v !== null)));
      if (isValid) {
        setKernel(parsedKernel as number[][][]);
      }
      return;
    }

    setSelectedPreset(CUSTOM_2D_OPTION);
    const nextDraft = draftGrayKernel.map((r) => [...r]);
    nextDraft[row][col] = value;
    setDraftGrayKernel(nextDraft);

    const parsedKernel = nextDraft.map((r) => r.map(parseCell));
    const isValid = parsedKernel.every((r) => r.every((v) => v !== null));
    if (isValid) {
      setKernel(parsedKernel as number[][]);
    }
  }

  function handleSizeChange(newWidth: number, newHeight: number) {
    const clampedWidth = Math.round(Math.min(MAX_KERNEL_SIZE, Math.max(MIN_KERNEL_SIZE, newWidth)));
    const clampedHeight = Math.round(Math.min(MAX_KERNEL_SIZE, Math.max(MIN_KERNEL_SIZE, newHeight)));

    function resizeMatrix(matrix: string[][], width: number, height: number): string[][] {
      return Array.from({ length: height }, (_, rowIndex) =>
        Array.from({ length: width }, (_, colIndex) => matrix[rowIndex]?.[colIndex] ?? "0"),
      );
    }

    if (Array.isArray(kernel[0][0])) {
      setSelectedPreset(CUSTOM_3D_OPTION);
      const resizedDraft = draftColorKernel.map((channelMatrix) =>
        resizeMatrix(channelMatrix, clampedWidth, clampedHeight),
      );
      setDraftColorKernel(resizedDraft);
      setKernel(
        resizedDraft.map((channelMatrix) =>
          channelMatrix.map((row) => row.map((cell) => parseCell(cell) ?? 0)),
        ),
      );
      return;
    }

    setSelectedPreset(CUSTOM_2D_OPTION);
    const resizedDraft = resizeMatrix(draftGrayKernel, clampedWidth, clampedHeight);
    setDraftGrayKernel(resizedDraft);
    setKernel(resizedDraft.map((row) => row.map((cell) => parseCell(cell) ?? 0)));
  }

  const channelLabels = ["R", "G", "B"];
  const is3DKernel = Array.isArray(kernel[0][0]);
  const activeMatrix = is3DKernel ? draftColorKernel[selectedChannel] : draftGrayKernel;
  const currentHeight = activeMatrix.length;
  const currentWidth = activeMatrix[0]?.length ?? 0;
  const isCustomKernel = selectedPreset === CUSTOM_2D_OPTION || selectedPreset === CUSTOM_3D_OPTION;

  return (
    <div className="flex flex-col w-full min-h-0 gap-3">
      <div className="grid grid-cols-[4.5rem_minmax(0,1fr)] items-center gap-2">
        <span className="text-sm font-bold text-gray-700 text-right">Kernel</span>
        <div className="justify-self-start [&_label]:hidden [&_select]:px-2 [&_select]:py-1 [&_select]:text-xs [&_optgroup]:font-bold [&_option]:font-normal">
          <Dropdown
            label=""
            groupedOptions={presetGroups}
            activeOption={selectedPreset}
            onChange={handlePresetChange}
          />
        </div>
      </div>
      { is3DKernel && (
        <div className="grid grid-cols-[4.5rem_minmax(0,1fr)] items-center gap-2">
          <span className="text-sm font-bold text-gray-700 text-right">Channel</span>
          <div className="flex gap-2 justify-self-start">
            {channelLabels.map((label, index) => (
              <button
                key={label}
                type="button"
                className={`px-3 py-1 text-sm rounded border ${
                  selectedChannel === index ? "bg-orange-200 border-orange-300" : "bg-white border-gray-300"
                }`}
                onClick={() => setSelectedChannel(index)}
              >
                {label}
              </button>
            ))}
          </div>
        </div>
      )}
      <div className="grid grid-cols-[4.5rem_minmax(0,1fr)] items-center gap-2">
        <span className="text-sm font-bold text-gray-700 text-right">Size</span>
        <div className="flex items-center gap-2 justify-self-start">
          <input
            type="number"
            min={MIN_KERNEL_SIZE}
            max={MAX_KERNEL_SIZE}
            disabled={!isCustomKernel}
            aria-label="Kernel width"
            className="w-20 px-2 py-1 text-sm text-center rounded border border-gray-300 disabled:bg-gray-100 disabled:text-gray-500"
            value={currentWidth}
            onChange={(event) => {
              if (Number.isFinite(event.currentTarget.valueAsNumber)) {
                handleSizeChange(event.currentTarget.valueAsNumber, currentHeight);
              }
            }}
          />
          <span className="text-sm">×</span>
          <input
            type="number"
            min={MIN_KERNEL_SIZE}
            max={MAX_KERNEL_SIZE}
            disabled={!isCustomKernel}
            aria-label="Kernel height"
            className="w-20 px-2 py-1 text-sm text-center rounded border border-gray-300 disabled:bg-gray-100 disabled:text-gray-500"
            value={currentHeight}
            onChange={(event) => {
              if (Number.isFinite(event.currentTarget.valueAsNumber)) {
                handleSizeChange(currentWidth, event.currentTarget.valueAsNumber);
              }
            }}
          />
        </div>
      </div>
      <div className="flex flex-col gap-2 ml-[5rem]">
        {activeMatrix.map((row, rowIndex) => (
          <div key={rowIndex} className="flex gap-2">
            {row.map((cellValue, colIndex) => (
              <input
                key={`${rowIndex}-${colIndex}`}
                type="text"
                disabled={!isCustomKernel}
                className="w-14 px-2 py-1 border border-gray-300 rounded text-sm disabled:bg-gray-100 disabled:text-gray-500"
                value={cellValue}
                onChange={(event) =>
                  handleKernelChange(
                    is3DKernel ? selectedChannel : 0,
                    rowIndex,
                    colIndex,
                    event.target.value,
                  )
                }
              />
            ))}
          </div>
        ))}
      </div>
      <NumericOption label="Stride" value={stride} min={1} onChange={setStride} />
      <NumericOption label="Dilation" value={dilation} min={1} onChange={setDilation} />
      <NumericOption label="Padding" value={padding} min={0} onChange={setPadding} />
    </div>
  );
}

function NumericOption({ label, value, min, onChange }: { label: string; value: number; min: number; onChange: (value: number) => void }) {
  return (
    <div className="grid grid-cols-[4.5rem_minmax(0,1fr)] items-center gap-2">
      <span className="text-sm font-bold text-gray-700 text-right">{label}</span>
      <input
        type="number"
        min={min}
        max={MAX_KERNEL_SIZE}
        aria-label={label}
        className="w-20 px-2 py-1 text-sm text-center rounded border border-gray-300"
        value={value}
        onChange={(event) => {
          const nextValue = event.currentTarget.valueAsNumber;
          if (Number.isFinite(nextValue)) {
            onChange(Math.round(Math.min(MAX_KERNEL_SIZE, Math.max(min, nextValue))));
          }
        }}
      />
    </div>
  );
}