import {
  Avatar as RadixAvatar,
  type AvatarProps as RadixAvatarProps,
} from '@radix-ui/themes';
import type { Responsive } from '@radix-ui/themes/props';
import { useMemo, forwardRef, useCallback } from 'react';

type Size = 'small' | 'medium' | 'large' | 'full';
type OriginalSize = '1' | '2' | '3' | '4' | '5' | '6' | '7' | '8' | '9';

export type AvatarProps = Omit<RadixAvatarProps, 'size'> & {
  size?: Responsive<Size>;
};

export const Avatar = forwardRef<HTMLImageElement, AvatarProps>(
  (props: AvatarProps, ref): React.ReactNode => {
    const { children, size, variant = 'soft', ...rest } = props;
    const getOriginalSize = useCallback((value?: Size): OriginalSize => {
      switch (value) {
        case 'small':
          return '1';
        case 'medium':
          return '2';
        case 'large':
          return '3';
        case 'full':
          return '4';
        default:
          return '3';
      }
    }, []);

    const radixSize = useMemo<RadixAvatarProps['size']>(() => {
      if (typeof size === 'string') {
        return getOriginalSize(size);
      }

      if (typeof size === 'object') {
        const map: RadixAvatarProps['size'] = {};
        let key: keyof typeof size;
        for (key in size) {
          map[key] = size[key] && getOriginalSize(size[key]);
        }
        return map;
      }
    }, [getOriginalSize, size]);

    return (
      <RadixAvatar {...rest} ref={ref} size={radixSize} variant={variant}>
        {children}
      </RadixAvatar>
    );
  }
);

Avatar.displayName = 'Avatar';
