Height state value not updating immediately and state value always one step behind

Question

I have a small project using justified-layout and react-virtualized,

I set the width and height of the preview images with the useCallback of previewImageContainerRef Ref. ( In the <PreviewImage> component [(<Results> / <PreviewImgRow> / <PreviewImage>) components > albums > results > preview_image.js], ref={previewImageContainerRef} is in preview_image.js as well, located on <div> with className="c-flex-center picture-frame-wrapper")

    // useCallback > Ref : previewImageContainerRef
    const previewImageContainerRef = useCallback(node => {

        if (node !== null) {

            // Set > Preview Image Width
            props.setPreviewImageWidth(node.getBoundingClientRect().width);

            // Calculate > Preview Image Height
            let new_height = (node.getBoundingClientRect().width / defaultImageWidthRef.current ) * defaultImageHeightRef.current;

            // Set > Preview Image Height
            props.setPreviewImageHeight(new_height);

        }

    }, []);

The useState for setPreviewImageWidth and setPreviewImageHeight, are defined in the parent component (<Results> [components > albums > results > Results.js]):

  // Preview Image > Width & Height
  const [previewImageWidth, setPreviewImageWidth] = useState(0);
  const [previewImageHeight, setPreviewImageHeight] = useState(0);

justifiedRowDetailsRef is defined in <Results> component, it uses useCallback as well and it is one of the parent <div> wrapping around the preview image (ref={justifiedRowDetailsRef} is in component [(<Results> / <PreviewImgRow>) components > albums > results > PreviewImgRow.js]):

  // useCallback > Ref : justifiedRowDetailsRef
  const justifiedRowDetailsRef = useCallback((node) => {
    if (node !== null) {

      setJustifiedRowDetailsWidth(node.getBoundingClientRect().width);
      setJustifiedRowDetailsHeight(node.getBoundingClientRect().height);

    }
  }, []);

The useState for justifiedRowDetailsWidth and justifiedRowDetailsHeight are defined in <Result> component:

  const [justifiedRowDetailsWidth, setJustifiedRowDetailsWidth] = useState(0);
  const [justifiedRowDetailsHeight, setJustifiedRowDetailsHeight] = useState(0);

justifiedRowDetailsHeight is used in the height calculations for multiple parent <div> wrapping the preview image

The problem I am facing is that the justifiedRowDetailsHeight value is only updated correctly when the image in the <JustifiedBox> component is clicked to expand the second time and after

Preview of the problem:

and here is my codesandbox link: https://codesandbox.io/s/gifted-ride-qvquwv?file=/components/albums/results/Results.js

any asistance on this is greatly appreciated

Edit: I have commented out:

  // useCallback > Ref : justifiedRowDetailsRef
  const justifiedRowDetailsRef = useCallback((node) => {
    if (node !== null) {

      setJustifiedRowDetailsWidth(node.getBoundingClientRect().width);
      setJustifiedRowDetailsHeight(node.getBoundingClientRect().height);

    }
  }, []);

and replaced it with:

const justifiedRowDetailsRef = useRef(null);

and passed setJustifiedRowDetailsWidth and setJustifiedRowDetailsHeight to <PreviewImgRow> component, and added this useEffect() into it:

  useEffect(() => {
    if (justifiedRowDetailsRef.current) {
      console.log("justifiedRowDetailsRef.current > NOT EMPTY");

      setJustifiedRowDetailsWidth(justifiedRowDetailsRef.current.offsetWidth);
      setJustifiedRowDetailsHeight(justifiedRowDetailsRef.current.offsetHeight);
    }
  }, [previewImageHeight, isPreviewShown]);

Now on first click of the image, the height is correct, but if I clicked on images with different dimension ( for example, the square image keep calm and eat food), the height of the preview image is still using the previous preview image height ( the rectangle height ) and one step behind

Steps to reproduce the bug:

Step1:

Step2:

Step3:

Step4: