Retention Curve

"use client";

import * as React from "react";
import { Area, AreaChart, CartesianGrid, XAxis, YAxis } from "recharts";

import { cn } from "@/lib/utils";
import { ChartCard } from "@/registry/default/widget-kit/chart-card";
import {
  ChartContainer,
  ChartTooltip,
  ChartTooltipContent,
  type ChartConfig,
} from "@/components/ui/chart";
import {
  resolveNarrative,
  type NarrativeProp,
} from "@/registry/default/chart-narrative/chart-narrative";
import { cohortStats } from "@/registry/default/sextant-insights/stats";
import type { Insight } from "@/registry/default/sextant-insights/types";
import { StudioTarget, useStudio } from "@/registry/default/widget-kit/studio";
import type { SizeProp } from "@/registry/default/widget-kit/chart-size";
import type {
  ChartClassNames,
  SeriesConfig,
} from "@/registry/default/widget-kit/contract";
import { safeData } from "@/registry/default/sextant-contracts/contracts";
import { ChartStates } from "@/registry/default/chart-states/chart-states";
import { retentionCohortsSchema, retentionCohorts, type CohortRow } from "./data";

// Re-export the data contract so consumers get the schema + types from a
// single import (no reaching into ./data).
export { retentionCohortsSchema, type CohortRow } from "./data";

/** The one encoding this widget draws: the average-retention curve. */
export type RetentionCurveSeries = { retention?: SeriesConfig };

type Props = {
  // ── Data ────────────────────────────────────────────────────────────────
  data?: CohortRow[];

  // ── Card chrome ───────────────────────────────────────────────────────────
  label?: string;
  attribution?: string;
  /** Shorthand for `classNames.root` — width, ring, shadow, margins, etc. */
  className?: string;
  /** Targeted class overrides per region (root/header/body/narrative/footnote). */
  classNames?: ChartClassNames;
  /** Force a size tier, or "auto" (default) to adapt to the card's own width. */
  size?: SizeProp;

  // ── Customize ───────────────────────────────────────────────────────────
  /** Recolor / rename / hide the curve: `series={{ retention: { color: "var(--brand)" } }}`. */
  series?: RetentionCurveSeries;
  /** Override percentage formatting (axis, hero, tooltip). */
  valueFormatter?: (n: number) => string;
  /** Replace, hide (`false`), or fully render the one-line narrative. */
  narrative?: NarrativeProp;
};

const fmtPct = (n: number): string => `${Math.round(n * 100)}%`;

export function ProductRetentionCurve(props: Props) {
  // Validate consumer data at the boundary: a bad shape shows the error
  // state instead of crashing. The built-in demo data is trusted, so skip it.
  const parsed =
    props.data === undefined ? null : safeData(retentionCohortsSchema, props.data);
  if (parsed && !parsed.ok) {
    return (
      <ChartStates status="error" shape="area" error={{ message: parsed.error }}>
        {null}
      </ChartStates>
    );
  }
  return <ProductRetentionCurveImpl {...props} />;
}

function ProductRetentionCurveImpl({
  data = retentionCohorts,
  label,
  attribution,
  className,
  classNames,
  size = "auto",
  series,
  valueFormatter,
  narrative,
}: Props) {
  const studio = useStudio();
  const title = studio.text("title", label ?? "Product retention curves", {
    label: "Chart title",
  });
  const gradientId = `rc-fill-${React.useId().replace(/:/g, "")}`;

  const color = series?.retention?.color ?? "var(--chart-1)";
  const seriesLabel = series?.retention?.label ?? "Avg retention";
  const fmt = valueFormatter ?? fmtPct;

  const stats = React.useMemo(() => cohortStats(data.map((r) => r.retention)), [data]);

  const chartData = React.useMemo(
    () => stats.byWeekMean.map((v, w) => ({ week: `W${w}`, retention: v })),
    [stats],
  );

  const lastWeek = Math.max(0, stats.byWeekMean.length - 1);
  const lastRetention = stats.byWeekMean[lastWeek] ?? 0;

  // Best post-week-0 cohort drives the narrative, compared to that week's mean.
  const best = stats.bestCohort;
  const insight: Insight | undefined =
    data.length > 0 && Number.isFinite(best.value)
      ? {
          kind: "cohort",
          cohortLabel: data[best.cohort]?.cohort ?? `#${best.cohort + 1}`,
          weekIndex: best.week,
          retention: best.value,
          cohortBaseline: stats.byWeekMean[best.week] ?? stats.overallMean,
        }
      : undefined;

  const config = {
    retention: { label: seriesLabel, color },
  } satisfies ChartConfig;

  return (
    <ChartCard.Root size={size} className={cn(className, classNames?.root)}>
      <ChartCard.Header className={classNames?.header}>
        <ChartCard.HeaderMain>
          <StudioTarget fieldKey="title" variant="inline">
            <ChartCard.Label>{title}</ChartCard.Label>
          </StudioTarget>
          <ChartCard.HeroNumber>{fmt(lastRetention)}</ChartCard.HeroNumber>
          <p className="tabular text-sm text-muted-foreground">
            Week {lastWeek} retention · {data.length} cohorts
          </p>
        </ChartCard.HeaderMain>
      </ChartCard.Header>

      {resolveNarrative(narrative, insight, { className: classNames?.narrative })}

      <ChartCard.Body className={classNames?.body}>
        <ChartContainer config={config} className="h-56 w-full">
          <AreaChart data={chartData} margin={{ left: 0, right: 8, top: 8 }}>
            <defs>
              <linearGradient id={gradientId} x1="0" y1="0" x2="0" y2="1">
                <stop offset="5%" stopColor={color} stopOpacity={0.3} />
                <stop offset="95%" stopColor={color} stopOpacity={0.02} />
              </linearGradient>
            </defs>
            <CartesianGrid vertical={false} stroke="var(--chart-grid)" />
            <XAxis
              dataKey="week"
              tickLine={false}
              axisLine={false}
              tickMargin={10}
              fontSize={10}
              className="smallcaps"
            />
            <YAxis
              tickFormatter={(v: number) => fmt(v)}
              tickLine={false}
              axisLine={false}
              width={40}
              domain={[0, 1]}
              fontSize={10}
              className="smallcaps tabular"
            />
            <ChartTooltip
              content={
                <ChartTooltipContent
                  labelFormatter={(label) => (
                    <span className="smallcaps-ink">{String(label)}</span>
                  )}
                  formatter={(value, name, item) => (
                    <>
                      <span
                        aria-hidden
                        className="h-2.5 w-2.5 shrink-0 rounded-[2px]"
                        style={{ backgroundColor: (item as { color?: string }).color }}
                      />
                      <span className="flex flex-1 items-center justify-between gap-3 leading-none">
                        <span className="text-muted-foreground">
                          {config[name as keyof typeof config]?.label ?? String(name)}
                        </span>
                        <span className="font-mono font-medium text-foreground tabular-nums">
                          {fmt(Number(value))}
                        </span>
                      </span>
                    </>
                  )}
                />
              }
            />
            <Area
              dataKey="retention"
              type="monotone"
              stroke={color}
              strokeWidth={1.5}
              fill={`url(#${gradientId})`}
            />
          </AreaChart>
        </ChartContainer>
      </ChartCard.Body>

      <ChartCard.Footnote className={classNames?.footnote}>
        {attribution ?? "Analytics engine · updated hourly"}
      </ChartCard.Footnote>
    </ChartCard.Root>
  );
}