/// <reference types="vite/client" />
import type { ComponentType } from 'react';
import React, { useState, useRef, useEffect } from 'react';

export function createComponentHmrWrapper<P extends Record<string, unknown>>(
  OriginalComponent: ComponentType<P>,
  hotData: {
    Component?: ComponentType<P>;
    setComponent?: (NewComponent: ComponentType<P>) => void;
  }
): ComponentType<P> {
  return import.meta.hot
    ? (props) => {
        const [WrappedComponent, setComponent] = useState(
          () => hotData.Component ?? OriginalComponent
        );
        const isMountedRef = useRef(false);

        useEffect(() => {
          hotData.setComponent = (NewComponent) => {
            hotData.Component = NewComponent;
            if (isMountedRef.current) {
              setComponent(() => NewComponent);
            }
          };
        }, [setComponent]);

        useEffect(() => {
          isMountedRef.current = true;

          return () => {
            isMountedRef.current = false;
          };
        }, []);
        return <WrappedComponent {...props} />;
      }
    : OriginalComponent;
}

type CustomElementClass = typeof HTMLElement;

const replaceMethodsOnElement = (
  elementNode: Element,
  newClass: CustomElementClass
) => {
  const instanceDescriptors = Object.getOwnPropertyDescriptors(newClass);
  for (const [name, descriptor] of Object.entries(instanceDescriptors)) {
    if (name !== 'length' && name !== 'name' && name !== 'prototype') {
      Object.defineProperty(elementNode.constructor, name, descriptor);
    }
  }

  const instanceSymbols = Object.getOwnPropertySymbols(newClass);
  for (const symbol of instanceSymbols) {
    //@ts-expect-error using symbol as index
    Object.defineProperty(elementNode, symbol, instanceDescriptors[symbol]);
  }

  const prototypeDescriptors = Object.getOwnPropertyDescriptors(
    newClass.prototype
  );
  for (const [name, descriptor] of Object.entries(prototypeDescriptors)) {
    Object.defineProperty(elementNode, name, descriptor);
  }

  const prototypeSymbols = Object.getOwnPropertySymbols(newClass.prototype);
  for (const symbol of prototypeSymbols) {
    //@ts-expect-error using symbol as index
    Object.defineProperty(elementNode, symbol, prototypeDescriptors[symbol]);
  }
};

function visitPrototypeChain(
  rootClass: CustomElementClass,
  visitFn: (chainClass: CustomElementClass) => void
) {
  const inheritanceArray = [rootClass];
  let currentClass = Object.getPrototypeOf(rootClass);

  while (currentClass !== HTMLElement) {
    inheritanceArray.unshift(currentClass);
    currentClass = Object.getPrototypeOf(currentClass);
  }

  for (const classForVisiting of inheritanceArray) {
    visitFn(classForVisiting);
  }
}

export function applyCustomElementHmr(
  elementNode: Element,
  NewElementClass: CustomElementClass
) {
  visitPrototypeChain(NewElementClass, (classInheritance) => {
    replaceMethodsOnElement(elementNode, classInheritance);
  });
}