import { Fragment } from 'react';

import { Box, Stack, Typography } from '@mui/material';
import { type Meta, type StoryObj } from '@storybook/react-vite';

import { type AIIconButtonProps } from './types';
import { AIIconButton } from '.';

const VARIANTS = ['primary', 'secondary', 'tertiary'] as const;

const STATES: { label: string; props: Partial<AIIconButtonProps> }[] = [
  { label: 'Enabled', props: {} },
  { label: 'Loading', props: { loading: true } },
  { label: 'Disabled', props: { disabled: true } },
];

const meta: Meta<typeof AIIconButton> = {
  title: 'Design System/AI/Buttons/AIIconButton',
  component: AIIconButton,
  tags: ['autodocs'],
  args: {
    size: 'large',
    icon: 'sparkles',
    variant: 'primary',
    'aria-label': 'AI action',
  },
  argTypes: {
    variant: {
      control: 'radio',
      options: VARIANTS,
      table: { defaultValue: { summary: 'primary' } },
    },
    size: {
      control: 'radio',
      options: ['small', 'large'],
      table: { defaultValue: { summary: 'large' } },
    },
    icon: {
      control: 'radio',
      options: ['sparkles', 'ai-pencil', 'ai-photo'],
      table: { defaultValue: { summary: 'sparkles' } },
    },
    loading: {
      control: 'boolean',
      description: 'Shows the progress/stop state',
    },
    onStop: { table: { disable: true } },
  },
};

export default meta;

type Story = StoryObj<typeof meta>;

const ColumnHeader = ({ label }: { label: string }) => (
  <Typography
    variant="caption"
    sx={{ fontWeight: 700, textAlign: 'center', color: 'text.secondary' }}
  >
    {label}
  </Typography>
);

const RowHeader = ({ label }: { label: string }) => (
  <Typography
    variant="subtitle2"
    sx={{ textTransform: 'capitalize' }}
  >
    {label}
  </Typography>
);

/** Type rows × state columns. Hover a button to see the gradient/sparkle animation. */
const StateMatrix = (args: AIIconButtonProps) => (
  <Box
    sx={{
      display: 'inline-grid',
      gridTemplateColumns: `auto repeat(${STATES.length}, 1fr)`,
      gap: 3,
      alignItems: 'center',
    }}
  >
    <span />
    {STATES.map(({ label }) => (
      <ColumnHeader
        key={label}
        label={label}
      />
    ))}

    {VARIANTS.map(variant => (
      <Fragment key={variant}>
        <RowHeader label={variant} />
        {STATES.map(({ label, props }) => (
          <Box
            key={label}
            sx={{ display: 'flex', justifyContent: 'center' }}
          >
            <AIIconButton
              {...args}
              variant={variant}
              {...props}
            />
          </Box>
        ))}
      </Fragment>
    ))}
  </Box>
);

export const Playground: Story = {
  args: {},
};

export const Overview: Story = {
  parameters: { controls: { disable: true } },
  render: StateMatrix,
};

export const Sizes: Story = {
  parameters: { controls: { disable: true } },
  render: (args: AIIconButtonProps) => (
    <Stack spacing={3}>
      {(['large', 'small'] as const).map(size => (
        <Stack
          key={size}
          spacing={1}
        >
          <RowHeader label={size} />
          <Stack
            direction="row"
            spacing={2}
            alignItems="center"
            flexWrap="wrap"
          >
            {VARIANTS.map(variant => (
              <AIIconButton
                key={variant}
                {...args}
                variant={variant}
                size={size}
              />
            ))}
          </Stack>
        </Stack>
      ))}
    </Stack>
  ),
};

export const Dark: Story = {
  parameters: { controls: { disable: true } },
  render: StateMatrix,
  globals: { themeMode: 'dark' },
};