CohortComparison

"use client";

import * as React from "react";
import {
  TabsRoot,
  TabList,
  Tab,
  TabPanel,
  StatCard,
  FilterBar,
  DataTable,
  Badge,
  Button,
  SurfaceLayout,
  TimeRangeSelector,
  type ColumnDef,
  type FilterGroupDef,
  type DateRange,
} from "@reopt-ai/opt-ui";
import {
  ComparisonChart,
  RetentionHeatmap,
  FunnelChart,
  type ChartDataPoint,
  type ComparisonSeriesDef,
  type RetentionRow,
  type FunnelDataPoint,
} from "@reopt-ai/opt-charts";

interface SummaryRow {
  id: string;
  [k: string]: unknown;
}

/** Card data for `CohortMetric`. */
export interface CohortMetricCard {
  id: string;
  cohort: string;
  title: string;
  value: string;
  change: string;
  trend: "up" | "down" | "neutral";
}

/** N-cohort definition for multi-segment comparison */
export interface CohortDef {
  id: string;
  label: string;
  color?: string;
  retentionData?: RetentionRow[];
  funnelData?: FunnelDataPoint[];
}

const DEFAULT_COHORT_COLORS = [
  "#3b82f6",
  "#8b5cf6",
  "#f59e0b",
  "#10b981",
  "#ef4444",
];

function createUniqueCohortId(
  cohorts: CohortDef[],
  nextIdRef: React.MutableRefObject<number>,
) {
  const reservedIds = new Set(cohorts.map((cohort) => cohort.id));
  let id = "";

  do {
    id = `cohort-${nextIdRef.current++}`;
  } while (reservedIds.has(id));

  return id;
}

/** Labels for `CohortComparison`. */
export interface CohortComparisonLabels {
  title?: string;
  overview?: string;
  retentionA?: string;
  retentionB?: string;
  funnelTab?: string;
  metricsTab?: string;
  summary?: string;
  cohortSelectorLabel?: string;
  filtersLabel?: string;
  detailsLabel?: string;
  noComparisonData?: string;
  emptyTitle?: string;
  emptyDescription?: string;
  /** Label for the add cohort button (N-cohort mode) */
  addCohortButton?: string;
  /** Label shown when max cohorts reached */
  maxCohortsReached?: string;
  /** Label for remove cohort action */
  removeCohort?: string;
}

const defaultLabels: Required<CohortComparisonLabels> = {
  title: "Cohort Comparison",
  overview: "Comparison",
  retentionA: "Cohort A",
  retentionB: "Cohort B",
  funnelTab: "Funnel",
  metricsTab: "Metrics",
  summary: "Summary",
  cohortSelectorLabel: "Cohort selector",
  filtersLabel: "Filters",
  detailsLabel: "Cohort comparison details",
  noComparisonData: "No comparison data",
  emptyTitle: "No cohort data",
  emptyDescription: "There is no cohort data to compare yet.",
  addCohortButton: "+ Add Cohort",
  maxCohortsReached: "Maximum cohorts reached",
  removeCohort: "Remove",
};

/** Props for `CohortComparison`. */
export interface CohortComparisonProps {
  comparisonData?: ChartDataPoint[];
  comparisonSeries?: ComparisonSeriesDef[];
  retentionDataA?: RetentionRow[];
  retentionDataB?: RetentionRow[];
  retentionPeriods?: string[];
  /** Funnel data per cohort for funnel comparison tab */
  funnelDataA?: FunnelDataPoint[];
  funnelDataB?: FunnelDataPoint[];
  /** N-cohort definitions (overrides A/B when provided) */
  cohorts?: CohortDef[];
  /** Callback when cohort list changes (add/remove) */
  onCohortsChange?: (cohorts: Array<{ id: string; name: string }>) => void;
  /** Maximum number of cohorts allowed (default: 5) */
  maxCohorts?: number;
  /** Metrics cards for cohort-level KPI comparison */
  metrics?: CohortMetricCard[];
  summaryData?: SummaryRow[];
  summaryColumns?: ColumnDef<SummaryRow>[];
  cohortALabel?: string;
  cohortBLabel?: string;
  /** FilterBar definitions (e.g., events, date range, segments) */
  filters?: FilterGroupDef[];
  onFilterChange?: (
    filterId: string,
    value: string | string[] | boolean,
  ) => void;
  header?: React.ReactNode;
  actions?: React.ReactNode;
  labels?: CohortComparisonLabels;
  loading?: boolean;
  timeRange?: DateRange;
  onTimeRangeChange?: (range: DateRange) => void;
  className?: string;
}

/** Renders the `CohortComparison` component. */
export function CohortComparison({
  comparisonData = [],
  comparisonSeries = [],
  retentionDataA = [],
  retentionDataB = [],
  retentionPeriods,
  funnelDataA,
  funnelDataB,
  cohorts,
  onCohortsChange,
  maxCohorts = 5,
  metrics,
  summaryData = [],
  summaryColumns = [
    { id: "metric", header: "Metric", accessor: "metric" as keyof SummaryRow },
    {
      id: "cohortA",
      header: "Cohort A",
      accessor: "cohortA" as keyof SummaryRow,
    },
    {
      id: "cohortB",
      header: "Cohort B",
      accessor: "cohortB" as keyof SummaryRow,
    },
    { id: "diff", header: "Difference", accessor: "diff" as keyof SummaryRow },
  ],
  cohortALabel = "Cohort A",
  cohortBLabel = "Cohort B",
  filters,
  onFilterChange,
  header,
  actions,
  labels: customLabels,
  loading = false,
  timeRange,
  onTimeRangeChange,
  className,
}: CohortComparisonProps) {
  const labels = { ...defaultLabels, ...customLabels };
  const sectionId = React.useId();
  const nextCohortIdRef = React.useRef(1);

  // N-cohort mode when cohorts prop is provided
  const useNCohort = cohorts && cohorts.length > 0;
  const canAddCohort = useNCohort && cohorts.length < maxCohorts;

  const isEmpty = useNCohort
    ? comparisonData.length === 0 &&
      cohorts.every((c) => !c.retentionData || c.retentionData.length === 0)
    : comparisonData.length === 0 &&
      retentionDataA.length === 0 &&
      retentionDataB.length === 0;

  const hasFilters = filters && filters.length > 0;
  const hasFunnel = useNCohort
    ? cohorts.some((c) => c.funnelData && c.funnelData.length > 0)
    : funnelDataA &&
      funnelDataA.length > 0 &&
      funnelDataB &&
      funnelDataB.length > 0;
  const hasMetrics = metrics && metrics.length > 0;

  const handleAddCohort = () => {
    if (!useNCohort || !onCohortsChange || !canAddCohort) return;
    const nextIndex = cohorts.length + 1;
    const newCohort = {
      id: createUniqueCohortId(cohorts, nextCohortIdRef),
      name: `Cohort ${nextIndex}`,
    };
    onCohortsChange([
      ...cohorts.map((c) => ({ id: c.id, name: c.label })),
      newCohort,
    ]);
  };

  const handleRemoveCohort = (cohortId: string) => {
    if (!useNCohort || !onCohortsChange) return;
    onCohortsChange(
      cohorts
        .filter((c) => c.id !== cohortId)
        .map((c) => ({ id: c.id, name: c.label })),
    );
  };

  return (
    <SurfaceLayout loading={loading} className={className}>
      {header || actions ? (
        <div className="flex items-center justify-between">
          {header && <div>{header}</div>}
          <div className="gap-element flex items-center">
            <TimeRangeSelector value={timeRange} onChange={onTimeRangeChange} />
            {actions}
          </div>
        </div>
      ) : (
        <div className="flex items-center justify-between">
          <h2 className="text-text-primary text-lg font-semibold">
            {labels.title}
          </h2>
          <TimeRangeSelector value={timeRange} onChange={onTimeRangeChange} />
        </div>
      )}

      {/* N-cohort selector */}
      {useNCohort && onCohortsChange && (
        <section
          aria-labelledby={`${sectionId}-cohort-selector`}
          className="gap-element flex flex-wrap items-center"
        >
          <span id={`${sectionId}-cohort-selector`} className="sr-only">
            {labels.cohortSelectorLabel}
          </span>
          {cohorts.map((c, idx) => {
            const color =
              c.color ??
              DEFAULT_COHORT_COLORS[idx % DEFAULT_COHORT_COLORS.length];
            return (
              <span
                key={c.id}
                className="inline-flex items-center gap-1.5 rounded-full border px-2.5 py-1"
                style={{ borderColor: color }}
              >
                <span
                  className="inline-block h-2.5 w-2.5 rounded-full"
                  style={{ backgroundColor: color }}
                  aria-hidden="true"
                />
                <Badge size="sm" variant="default">
                  {c.label}
                </Badge>
                {cohorts.length > 1 && (
                  <button
                    type="button"
                    onClick={() => handleRemoveCohort(c.id)}
                    className="text-text-tertiary hover:text-text-primary ml-0.5 text-xs"
                    aria-label={`${labels.removeCohort} ${c.label}`}
                  >
                    &times;
                  </button>
                )}
              </span>
            );
          })}
          {canAddCohort ? (
            <Button variant="ghost" size="sm" onClick={handleAddCohort}>
              {labels.addCohortButton}
            </Button>
          ) : (
            cohorts.length >= maxCohorts && (
              <span className="text-text-tertiary text-xs">
                {labels.maxCohortsReached}
              </span>
            )
          )}
        </section>
      )}

      {hasFilters && (
        <section aria-labelledby={`${sectionId}-filters`}>
          <span id={`${sectionId}-filters`} className="sr-only">
            {labels.filtersLabel}
          </span>
          <FilterBar filters={filters} onFilterChange={onFilterChange} />
        </section>
      )}

      {isEmpty ? (
        <div
          role="status"
          className="flex flex-col items-center justify-center py-16 text-center"
        >
          <h3 className="text-text-primary text-lg font-medium">
            {labels.emptyTitle}
          </h3>
          <p className="text-text-tertiary mt-1 text-sm">
            {labels.emptyDescription}
          </p>
        </div>
      ) : (
        <section aria-labelledby={`${sectionId}-tabs`}>
          <span id={`${sectionId}-tabs`} className="sr-only">
            {labels.detailsLabel}
          </span>
          <TabsRoot defaultSelectedId="comparison">
            <TabList>
              <Tab id="comparison">{labels.overview}</Tab>
              {useNCohort ? (
                cohorts.map((c) => (
                  <Tab key={c.id} id={`retention-${c.id}`}>
                    {c.label}
                  </Tab>
                ))
              ) : (
                <>
                  <Tab id="retention-a">{cohortALabel}</Tab>
                  <Tab id="retention-b">{cohortBLabel}</Tab>
                </>
              )}
              {hasFunnel && <Tab id="funnel">{labels.funnelTab}</Tab>}
              {hasMetrics && <Tab id="metrics">{labels.metricsTab}</Tab>}
              <Tab id="summary">{labels.summary}</Tab>
            </TabList>
            <TabPanel tabId="comparison">
              <div className="pt-4">
                {comparisonData.length > 0 && comparisonSeries.length > 0 ? (
                  <ComparisonChart
                    data={comparisonData}
                    series={comparisonSeries}
                    aria-label="Cohort comparison"
                  />
                ) : (
                  <p className="text-text-tertiary py-8 text-center text-sm">
                    {labels.noComparisonData}
                  </p>
                )}
              </div>
            </TabPanel>
            {useNCohort ? (
              cohorts.map((c) => (
                <TabPanel key={c.id} tabId={`retention-${c.id}`}>
                  <div className="pt-4">
                    <RetentionHeatmap
                      data={c.retentionData ?? []}
                      periods={retentionPeriods}
                      aria-label={`${c.label} retention`}
                    />
                  </div>
                </TabPanel>
              ))
            ) : (
              <>
                <TabPanel tabId="retention-a">
                  <div className="pt-4">
                    <RetentionHeatmap
                      data={retentionDataA}
                      periods={retentionPeriods}
                      aria-label={`${cohortALabel} retention`}
                    />
                  </div>
                </TabPanel>
                <TabPanel tabId="retention-b">
                  <div className="pt-4">
                    <RetentionHeatmap
                      data={retentionDataB}
                      periods={retentionPeriods}
                      aria-label={`${cohortBLabel} retention`}
                    />
                  </div>
                </TabPanel>
              </>
            )}
            {hasFunnel && (
              <TabPanel tabId="funnel">
                <div className="gap-section grid grid-cols-1 pt-4 md:grid-cols-2">
                  {useNCohort ? (
                    cohorts
                      .filter((c) => c.funnelData && c.funnelData.length > 0)
                      .map((c) => (
                        <div key={c.id}>
                          <h4 className="text-text-secondary mb-2 text-sm font-medium">
                            {c.label}
                          </h4>
                          <FunnelChart
                            data={c.funnelData!}
                            aria-label={`${c.label} funnel`}
                          />
                        </div>
                      ))
                  ) : (
                    <>
                      <div>
                        <h4 className="text-text-secondary mb-2 text-sm font-medium">
                          {cohortALabel}
                        </h4>
                        <FunnelChart
                          data={funnelDataA!}
                          aria-label={`${cohortALabel} funnel`}
                        />
                      </div>
                      <div>
                        <h4 className="text-text-secondary mb-2 text-sm font-medium">
                          {cohortBLabel}
                        </h4>
                        <FunnelChart
                          data={funnelDataB!}
                          aria-label={`${cohortBLabel} funnel`}
                        />
                      </div>
                    </>
                  )}
                </div>
              </TabPanel>
            )}
            {hasMetrics && (
              <TabPanel tabId="metrics">
                <div className="gap-group grid grid-cols-2 pt-4 sm:grid-cols-3">
                  {metrics.map((m) => (
                    <div key={m.id} className="flex flex-col gap-1">
                      <span className="text-text-tertiary text-xs font-medium">
                        {m.cohort}
                      </span>
                      <StatCard
                        id={m.id}
                        title={m.title}
                        value={m.value}
                        change={m.change}
                        trend={m.trend}
                      />
                    </div>
                  ))}
                </div>
              </TabPanel>
            )}
            <TabPanel tabId="summary">
              <div className="pt-4">
                <DataTable
                  columns={summaryColumns}
                  data={summaryData}
                  keyExtractor={(r) => r.id}
                />
              </div>
            </TabPanel>
          </TabsRoot>
        </section>
      )}
    </SurfaceLayout>
  );
}

CohortComparison.displayName = "CohortComparison";