import clsx from 'clsx';
import React, { useMemo } from 'react';
import { twMerge } from 'tailwind-merge';

import Badge from '@/components/Badge/Badge';
import TitleText from '@/components/title/TitleText';
import { Tooltip } from '@/components/tooltip/Tooltip';
import { CheckIcon, InfoIcon, LockIcon } from '@/icons';
import {
  UsagePlanDescription,
  UsagePlanSchema,
  UsagePlanTableComparison,
} from '@/state/sessionStore';
import { renderIcon } from '@/utils/iconUtils';

import { BadgeConfig, getPlanBadge } from './badgeConfig';

interface HeaderCellProps {
  name: string;
  badge?: BadgeConfig;
  isLast?: boolean;
}

const HeaderCell: React.FC<HeaderCellProps> = ({ name, badge, isLast }) => (
  <div
    className={twMerge(
      'flex flex-row flex-wrap items-center justify-center gap-2 border-b border-border-primary p-4',
      !isLast && 'border-r border-border-primary'
    )}
  >
    <TitleText text={name} subtitle={true} />
    {badge && (
      <Badge text={badge.text} icon={badge.icon} className={badge.className} />
    )}
  </div>
);

interface CellProps {
  value: string | boolean;
  align?: 'start' | 'center';
  rowHeader?: boolean;
  isLast?: boolean;
  tooltip?: string | null;
}

const Cell: React.FC<CellProps> = ({
  value,
  align = 'center',
  rowHeader = false,
  isLast = false,
  tooltip,
}) => {
  return (
    <div
      className={twMerge(
        'flex items-center border-t border-border-primary p-4 text-sm',
        clsx(
          align === 'start' ? 'justify-start' : 'justify-center',
          !isLast && 'border-r border-border-primary',
          rowHeader
            ? 'font-medium text-foreground-primary'
            : 'text-foreground-secondary',
          rowHeader && tooltip ? 'justify-between' : ''
        )
      )}
    >
      <span>
        {typeof value === 'boolean' ? (
          value ? (
            <CheckIcon
              width={20}
              height={20}
              className='text-accent-success-on-primary'
            />
          ) : (
            <LockIcon
              width={20}
              height={20}
              className='text-foreground-tertiary'
            />
          )
        ) : (
          value
        )}
      </span>
      {rowHeader && tooltip && (
        <Tooltip label={tooltip}>
          <InfoIcon
            width={24}
            height={24}
            className='cursor-help text-foreground-secondary hover:text-foreground-primary'
          />
        </Tooltip>
      )}
    </div>
  );
};

export interface ComparePlansTableProps {
  tableComparison?: UsagePlanTableComparison | null;
  currentPlanKey?: string | null;
  plans?: UsagePlanSchema[];
  usagePlanDescriptions?: Record<string, UsagePlanDescription>;
}

export const ComparePlansTable = React.forwardRef<
  HTMLDivElement,
  ComparePlansTableProps
>(({ tableComparison, currentPlanKey, plans, usagePlanDescriptions }, ref) => {
  const data = tableComparison?.table_sections ?? [];

  const planLabelsWithBadges = useMemo(() => {
    if (!plans) {
      return [];
    }

    return [...plans]
      .sort((a, b) => a.level - b.level)
      .map((plan) => {
        const planKey: string = plan.plan_key;
        const isUserActivePlan = currentPlanKey === plan.plan_key;
        const planDescription = usagePlanDescriptions?.[planKey] ?? undefined;
        const badge = getPlanBadge(planDescription, isUserActivePlan);

        return {
          plan_key: planKey.toLowerCase(),
          name: plan.name,
          badge,
        };
      });
  }, [plans, currentPlanKey, usagePlanDescriptions]);

  // Calculate dynamic grid configuration
  const totalColumns = planLabelsWithBadges.length + 1; // +1 for the description column

  // If no data is available after processing, don't render the component
  if (data.length === 0 || planLabelsWithBadges.length === 0) {
    return null;
  }

  return (
    <div ref={ref} className='flex hidden w-full flex-col gap-6 md:flex'>
      <h2 className='text-center text-3xl font-medium'>Compare Suno plans</h2>

      {/* Main grid container */}
      <div
        className={`grid overflow-hidden border-t border-r border-b border-l border-border-primary`}
        style={{
          gridTemplateColumns: `1fr repeat(${planLabelsWithBadges.length}, minmax(0, 1fr))`,
        }}
      >
        {/* Header row */}
        <div className='border-r border-b border-border-primary p-4' />
        {planLabelsWithBadges.map(({ name, badge }, index) => (
          <HeaderCell
            key={index}
            name={name}
            badge={badge}
            isLast={index === planLabelsWithBadges.length - 1}
          />
        ))}

        {/* Content rows with subgrids */}
        {data.map((section, sectionIndex) => (
          <React.Fragment key={`section-${sectionIndex}`}>
            {/* Section header */}
            <div
              className='flex items-center gap-2 bg-foreground-primary/5 px-6 py-3 text-sm font-medium text-foreground-secondary'
              style={{ gridColumn: `1 / ${totalColumns + 1}` }}
            >
              {renderIcon(section.icon, { width: 16, height: 16 })}
              {section.name}
            </div>

            {/* Features for this section */}
            {section.features.map((feature, featureIndex) => (
              <React.Fragment key={`feature-${sectionIndex}-${featureIndex}`}>
                {/* Feature row as a subgrid to maintain column alignment */}
                <div
                  className='grid'
                  style={{
                    gridColumn: `1 / ${totalColumns + 1}`,
                    gridTemplateColumns: `1fr repeat(${planLabelsWithBadges.length}, minmax(0, 1fr))`,
                  }}
                >
                  <Cell
                    value={feature.name}
                    align='start'
                    rowHeader={true}
                    tooltip={feature.tooltip}
                  />
                  {/* Dynamically map cells based on the plan keys */}
                  {planLabelsWithBadges.map((plan, index) => (
                    <Cell
                      key={index}
                      value={
                        (feature[plan.plan_key] as boolean | string | null) ||
                        false
                      }
                      isLast={index === planLabelsWithBadges.length - 1}
                    />
                  ))}
                </div>
              </React.Fragment>
            ))}
          </React.Fragment>
        ))}
      </div>
    </div>
  );
});

ComparePlansTable.displayName = 'ComparePlansTable';