import { useState, useEffect, useRef, useCallback } from 'react';
import { useLocation } from 'react-router-dom';

import { enhancedSmoothstep } from '../utils';

type HeadingEntry = {
  [key: string]: IntersectionObserverEntry;
};

export function useActiveHeading(
  contentElement: HTMLDivElement | null,
  displayedHeadings?: Array<{ id: string; depth: number } | null>,
): { heading: string | undefined; handleHeadingClick: (headingId: string) => void } {
  const [heading, setHeading] = useState<string | undefined>(undefined);
  const [headingElements, setHeadingElements] = useState<HTMLElement[]>([]);
  const headingElementsRef = useRef<HeadingEntry>({});
  const displayedHeadingsRef = useRef(displayedHeadings);
  const lockObserver = useRef<boolean>(false);
  const location = useLocation();

  useEffect(() => {
    setHeading(undefined);
    headingElementsRef.current = {};
  }, [location.pathname]);

  useEffect(() => {
    displayedHeadingsRef.current = displayedHeadings;
  }, [displayedHeadings]);

  const getVisibleHeadings = () => {
    const visibleHeadings: IntersectionObserverEntry[] = [];

    for (const key in headingElementsRef.current) {
      const headingElement = headingElementsRef.current[key];

      if (headingElement.isIntersecting) {
        visibleHeadings.push(headingElement);
      }
    }

    return visibleHeadings;
  };

  const isTopOfPage = () => window.scrollY <= 50;

  // Returns viewport ratio for intersection observer based on screen height
  const getViewportRatio = () => {
    const vh = window.innerHeight;
    if (vh >= 1400) return 0.4;
    if (vh >= 1000) return 0.3;
    return 0.25;
  };

  const getIndexFromId = useCallback(
    (id: string) => {
      return headingElements.findIndex((item) => item.id === id);
    },
    [headingElements],
  );

  const findHeaders = (allContent: HTMLDivElement) => {
    const allHeaders = allContent.querySelectorAll<HTMLElement>('.heading-anchor');
    const headers = Array.from(allHeaders);

    if (displayedHeadingsRef.current && displayedHeadingsRef.current.length > 0) {
      const displayedIds = displayedHeadingsRef.current.map((h) => h?.id).filter(Boolean);
      return headers.filter((header) => displayedIds.includes(header.id));
    }

    return headers;
  };

  const calculateVirtualPositions = useCallback((headers: HTMLElement[]) => {
    if (!headers.length) return {};

    const navbarHeight = document.querySelector('nav')?.getBoundingClientRect().height || 0;
    const viewportHeight = window.innerHeight;
    const triggerOffset = navbarHeight + viewportHeight * 0.25;

    const lastHeader = headers[headers.length - 1];
    const lastHeaderTop = lastHeader.offsetTop;
    const maxScrollableHeight = document.body.scrollHeight - viewportHeight;

    const requiredUplift = Math.max(0, lastHeaderTop - maxScrollableHeight + triggerOffset);

    const virtualPositions: Record<string, number> = {};

    for (let i = 0; i < headers.length; i++) {
      const header = headers[i];
      const headerTop = header.offsetTop;

      const normalizedPosition = headers.length === 1 ? 0 : i / (headers.length - 1);

      const adjustmentFactor = enhancedSmoothstep(normalizedPosition, 0.4);

      const uplift = requiredUplift * adjustmentFactor;
      const virtualTop = headerTop - uplift;

      virtualPositions[header.id] = virtualTop;
    }

    return virtualPositions;
  }, []);

  const getTriggerOffset = useCallback(() => {
    const navbarHeight = document.querySelector('nav')?.getBoundingClientRect().height || 0;
    return navbarHeight + window.innerHeight * getViewportRatio();
  }, []);

  const findActiveHeaderByVirtualPosition = useCallback(
    (headers: HTMLElement[], scrollY: number, triggerOffset: number) => {
      const virtualPositions = calculateVirtualPositions(headers);

      let activeHeader = headers[0];
      for (const header of headers) {
        const virtualTop = virtualPositions[header.id] || header.offsetTop;
        if (virtualTop <= scrollY + triggerOffset) {
          activeHeader = header;
        } else {
          break;
        }
      }
      return activeHeader;
    },
    [calculateVirtualPositions],
  );

  const intersectionCallback = useCallback(
    (entries: IntersectionObserverEntry[]) => {
      if (lockObserver.current) {
        return;
      }

      headingElementsRef.current = entries.reduce(
        (map: HeadingEntry, entry: IntersectionObserverEntry) => {
          map[entry.target.id] = entry;
          return map;
        },
        headingElementsRef.current,
      );

      if (isTopOfPage()) {
        const newHeading = headingElements[0]?.id || undefined;
        setHeading(newHeading);
        return;
      }

      const visibleHeadings = getVisibleHeadings();
      if (!visibleHeadings.length) {
        const triggerOffset = getTriggerOffset();

        if (isTopOfPage()) {
          setHeading(headingElements[0]?.id || undefined);
          return;
        }

        const totalHeight = window.scrollY + window.innerHeight;
        if (totalHeight >= document.body.scrollHeight - 10) {
          setHeading(headingElements[headingElements.length - 1]?.id || undefined);
          return;
        }

        const activeHeader = findActiveHeaderByVirtualPosition(
          headingElements,
          window.scrollY,
          triggerOffset,
        );
        setHeading(activeHeader.id);
        return;
      }

      if (visibleHeadings.length === 1) {
        setHeading(visibleHeadings[0].target.id);
        return;
      }

      visibleHeadings.sort((a, b) => {
        return getIndexFromId(a.target.id) - getIndexFromId(b.target.id);
      });

      const firstVisibleHeading = visibleHeadings[0];

      if (getIndexFromId(firstVisibleHeading.target.id) === 0 && isTopOfPage()) {
        setHeading(firstVisibleHeading.target.id);
        return;
      }

      const totalHeight = scrollY + window.innerHeight;
      if (totalHeight >= document.body.scrollHeight - 10) {
        setHeading(visibleHeadings[visibleHeadings.length - 1].target.id);
        return;
      }

      setHeading(visibleHeadings[0].target.id);
    },
    [getIndexFromId, headingElements, getTriggerOffset, findActiveHeaderByVirtualPosition],
  );

  const handleHeadingClick = useCallback(
    (headingId: string) => {
      lockObserver.current = true;
      setHeading(headingId);

      const headingElement = headingElements.find((el) => el.id === headingId);
      if (headingElement && headingElement.scrollIntoView) {
        headingElement.scrollIntoView({ behavior: 'smooth', block: 'start' });
      }

      setTimeout(() => {
        headingElementsRef.current = {};
        lockObserver.current = false;
      }, 300);
    },
    [headingElements],
  );

  useEffect(() => {
    if (!contentElement) {
      return;
    }
    const headers = findHeaders(contentElement);
    setHeadingElements(headers);

    if (headers.length > 0 && !heading) {
      if (isTopOfPage()) {
        setHeading(headers[0].id);
      } else {
        const totalHeight = scrollY + window.innerHeight;
        if (totalHeight >= document.body.scrollHeight - 10) {
          setHeading(headers[headers.length - 1]?.id || undefined);
        } else {
          const triggerOffset = getTriggerOffset();
          const activeHeader = findActiveHeaderByVirtualPosition(headers, scrollY, triggerOffset);
          setHeading(activeHeader.id);
        }
      }
    }
  }, [
    contentElement,
    location,
    heading,
    calculateVirtualPositions,
    getTriggerOffset,
    findActiveHeaderByVirtualPosition,
  ]);

  useEffect(() => {
    if (!headingElements?.length) {
      return;
    }
    headingElementsRef.current = {};

    const triggerOffset = getTriggerOffset();

    const observer = new IntersectionObserver(intersectionCallback, {
      rootMargin: `-${triggerOffset}px 0px -${window.innerHeight * getViewportRatio()}px 0px`,
      threshold: [0, 0.1, 0.3, 0.5, 0.7, 1],
    });

    headingElements.forEach((element) => {
      observer.observe(element);
    });

    const handleScroll = () => {
      if (lockObserver.current) return;

      if (isTopOfPage()) {
        setHeading(headingElements[0]?.id || undefined);
      }
    };

    window.addEventListener('scroll', handleScroll, { passive: true });

    return () => {
      observer.disconnect();
      window.removeEventListener('scroll', handleScroll);
    };
  }, [headingElements, intersectionCallback, getTriggerOffset]);

  return { heading, handleHeadingClick };
}