Spend Saturation

"use client";

import * as React from "react";
import {
  CartesianGrid,
  ComposedChart,
  Line,
  ReferenceLine,
  Scatter,
  XAxis,
  YAxis,
} from "recharts";

import { cn } from "@/lib/utils";
import { ChartCard } from "@/registry/default/widget-kit/chart-card";
import { ChartContainer, ChartTooltip, type ChartConfig } from "@/components/ui/chart";
import {
  ChartNarrative,
  resolveNarrative,
  type NarrativeProp,
} from "@/registry/default/chart-narrative/chart-narrative";
import { StudioTarget, useStudio } from "@/registry/default/widget-kit/studio";
import { SparkTile } from "@/registry/default/widget-kit/spark-tile";
import { useSizePlan } from "@/registry/default/widget-kit/use-chart-size";
import type { SizeProp } from "@/registry/default/widget-kit/chart-size";
import type { SizePlan } from "@/registry/default/widget-kit/size-plan";
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 { spendSaturationSchema, spendSaturation, type SaturationPoint } from "./data";

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

const COLOR_PRESETS = [
  { name: "Maroon (brand)", value: "var(--chart-1)" },
  { name: "Ink", value: "var(--chart-2)" },
  { name: "Cool gray", value: "var(--chart-3)" },
  { name: "Sage", value: "var(--chart-positive)" },
  { name: "Terracotta", value: "var(--chart-negative)" },
];

const SAT_PLAN: SizePlan = {
  lg: { shape: "full", height: "h-64", legend: "full" },
  md: { shape: "full", height: "h-56", legend: "inline" },
  sm: { shape: "compact" },
  mobile: { shape: "full", height: "h-64", legend: "inline", stack: true },
};

/**
 * Two encodings: the fitted response curve (`fit`) and the observed weekly
 * points (`actual`). Recolor either: `series={{ fit: { color: "var(--brand)" } }}`.
 */
export type SpendSaturationSeries = {
  fit?: SeriesConfig;
  actual?: SeriesConfig;
};

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

  // ── 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 the fit curve / observed points. */
  series?: SpendSaturationSeries;
  /** Override conversion formatting (Y axis, tooltip). Default whole number. */
  valueFormatter?: (n: number) => string;
  /** Replace, hide (`false`), or fully render the one-line narrative. */
  narrative?: NarrativeProp;
};

const fmtUsd = (n: number): string =>
  n >= 1000 ? `$${(n / 1000).toFixed(0)}k` : `$${Math.round(n)}`;
const fmtConv = (n: number): string => `${Math.round(n)}`;

const SAMPLES = 20;

type Datum = {
  spend: number;
  fit: number;
  actual: number | null;
  /** Cost of the next conversion at this spend = 1 / (dConversions/dSpend). */
  marginalCost: number;
};

export type SaturationModel = {
  chartData: Datum[];
  knee: number;
  marginalCostAtMax: number;
  marginalCostAtMin: number;
};

/**
 * Pure: fit conversions = a·√spend + b by least squares on u = √spend, then
 * derive the smooth curve samples, the diminishing-returns knee, and the
 * marginal cost-per-conversion at the spend extremes. Exported so the headless
 * `<MarketingSpendSaturationPlot>` and your own code derive points identically.
 */
export function buildSaturationModel(data: SaturationPoint[]): SaturationModel {
  const pts = data.filter((d) => d.spend > 0);
  const u = pts.map((d) => Math.sqrt(d.spend));
  const y = pts.map((d) => d.conversions);
  const n = u.length;
  const su = u.reduce((s, v) => s + v, 0);
  const sy = y.reduce((s, v) => s + v, 0);
  const suu = u.reduce((s, v) => s + v * v, 0);
  const suy = u.reduce((s, v, i) => s + v * y[i], 0);
  const denom = n * suu - su * su;
  const a = denom !== 0 ? (n * suy - su * sy) / denom : 0;
  const b = n > 0 ? (sy - a * su) / n : 0;
  const spends = pts.map((d) => d.spend);
  const xmin = spends.length ? Math.min(...spends) : 0;
  const xmax = spends.length ? Math.max(...spends) : 0;

  const fit = (x: number) => a * Math.sqrt(Math.max(x, 1)) + b;
  const marginal = (x: number) => a / (2 * Math.sqrt(Math.max(x, 1)));

  // Knee: where marginal yield halves vs the smallest observed spend (4× it).
  const knee = Math.max(xmin, Math.min(xmax, 4 * xmin));
  const mMax = marginal(xmax);
  const mMin = marginal(xmin);

  // Marginal cost of the next conversion at spend x = 1 / (dConversions/dSpend).
  const costOf = (x: number) => {
    const m = marginal(x);
    return m > 0 ? 1 / m : 0;
  };

  const chartData: Datum[] = pts.map((d) => ({
    spend: d.spend,
    fit: fit(d.spend),
    actual: d.conversions,
    marginalCost: costOf(d.spend),
  }));
  if (xmax > xmin) {
    for (let i = 0; i <= SAMPLES; i++) {
      const x = Math.round(xmin + ((xmax - xmin) * i) / SAMPLES);
      chartData.push({ spend: x, fit: fit(x), actual: null, marginalCost: costOf(x) });
    }
  }
  chartData.sort((p, q) => p.spend - q.spend);

  return {
    chartData,
    knee,
    marginalCostAtMax: mMax > 0 ? 1 / mMax : 0,
    marginalCostAtMin: mMin > 0 ? 1 / mMin : 0,
  };
}

/**
 * Curated hover card: the spend, expected vs observed conversions, and — the
 * chart's whole point — the marginal cost of the next conversion at that spend.
 * Replaces the generic content, which dumped the raw x-value (`7100`) and the
 * fit as unlabeled bare numbers.
 */
function SaturationTooltip({
  active,
  payload,
  fmt,
  fitColor,
  actualColor,
}: {
  active?: boolean;
  payload?: { payload?: Datum }[];
  fmt: (n: number) => string;
  fitColor: string;
  actualColor: string;
}) {
  if (!active || !payload?.length) return null;
  const row = payload[0]?.payload;
  if (!row) return null;
  return (
    <div className="grid min-w-[9rem] gap-2 rounded-lg border border-border/50 bg-background px-3 py-2 text-xs shadow-xl">
      <div className="tabular text-sm font-semibold text-foreground">
        {fmtUsd(row.spend)}
        <span className="ml-1 text-xs font-normal text-muted-foreground">
          / wk spend
        </span>
      </div>
      <div className="grid gap-1">
        <SaturationTipRow
          color={fitColor}
          label="Expected"
          value={`~${fmt(row.fit)}`}
        />
        {row.actual != null && (
          <SaturationTipRow
            color={actualColor}
            label="Weekly actual"
            value={fmt(row.actual)}
          />
        )}
      </div>
      <div className="flex items-center justify-between gap-6 border-t border-border/50 pt-1.5">
        <span className="text-muted-foreground">Next conversion</span>
        <span className="tabular font-medium text-foreground">
          {fmtUsd(row.marginalCost)}
        </span>
      </div>
    </div>
  );
}

function SaturationTipRow({
  color,
  label,
  value,
}: {
  color: string;
  label: string;
  value: string;
}) {
  return (
    <div className="flex items-center justify-between gap-6">
      <span className="inline-flex items-center gap-1.5 text-muted-foreground">
        <span
          aria-hidden
          className="size-2 rounded-[2px]"
          style={{ backgroundColor: color }}
        />
        {label}
      </span>
      <span className="tabular font-medium text-foreground">{value}</span>
    </div>
  );
}

export function MarketingSpendSaturation(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(spendSaturationSchema, props.data);
  if (parsed && !parsed.ok) {
    return (
      <ChartStates status="error" shape="line" error={{ message: parsed.error }}>
        {null}
      </ChartStates>
    );
  }
  return <MarketingSpendSaturationImpl {...props} />;
}

function MarketingSpendSaturationImpl({
  data = spendSaturation,
  label,
  attribution,
  className,
  classNames,
  size = "auto",
  series,
  valueFormatter,
  narrative,
}: Props) {
  const studio = useStudio();
  const title = studio.text("title", label ?? "Paid acquisition saturation curve", {
    label: "Chart title",
  });
  // `series.*.color` sets the default; Studio (when mounted) still edits on top.
  const fitColor = studio.color("fitColor", series?.fit?.color ?? "var(--chart-1)", {
    label: "Curve color",
    presets: COLOR_PRESETS,
  });
  const actualColor = studio.color(
    "actualColor",
    series?.actual?.color ?? "var(--chart-3)",
    { label: "Points color", presets: COLOR_PRESETS },
  );
  const attributionText = studio.text(
    "attribution",
    attribution ??
      "Acme marketing · Q1 2026 spend vs qualified conversions · √ response fit",
    { label: "Attribution" },
  );

  const model = React.useMemo(() => buildSaturationModel(data), [data]);

  // No template kind expresses "$ per next conversion"; compose the sentence.
  const autoNarrative = `Returns flatten beyond ${fmtUsd(model.knee)}/wk — the next conversion costs ${fmtUsd(model.marginalCostAtMax)} vs ${fmtUsd(model.marginalCostAtMin)} at low spend.`;

  const config = {
    fit: { label: series?.fit?.label ?? "Response curve", color: fitColor },
    actual: { label: series?.actual?.label ?? "Weekly actuals", color: actualColor },
  } satisfies ChartConfig;

  return (
    <ChartCard.Root
      size={size}
      className={cn(className, classNames?.root)}
      exportable
      exportTitle={title}
    >
      <SaturationInner
        data={data}
        config={config}
        model={model}
        title={title}
        autoNarrative={autoNarrative}
        narrative={narrative}
        actualColor={actualColor}
        fitColor={fitColor}
        valueFormatter={valueFormatter}
        attributionText={attributionText}
        classNames={classNames}
      />
    </ChartCard.Root>
  );
}

type InnerProps = {
  data: SaturationPoint[];
  config: ChartConfig;
  model: SaturationModel;
  title: string;
  autoNarrative: string;
  narrative: NarrativeProp | undefined;
  actualColor: string;
  fitColor: string;
  valueFormatter?: (n: number) => string;
  attributionText: string;
  classNames?: ChartClassNames;
};

function SaturationInner({
  data,
  config,
  model,
  title,
  autoNarrative,
  narrative,
  actualColor,
  fitColor,
  valueFormatter,
  attributionText,
  classNames,
}: InnerProps) {
  const spec = useSizePlan(SAT_PLAN);

  if (spec.shape === "compact") {
    return (
      <SparkTile
        label={title}
        value={fmtUsd(model.marginalCostAtMax)}
        chip={
          narrative === false ? undefined : (
            <ChartNarrative
              text={typeof narrative === "string" ? narrative : autoNarrative}
              variant="chip"
              className={classNames?.narrative}
            />
          )
        }
      />
    );
  }

  return (
    <>
      <ChartCard.Header className={classNames?.header}>
        <ChartCard.HeaderMain>
          <StudioTarget fieldKey="title" variant="inline">
            <ChartCard.Label>{title}</ChartCard.Label>
          </StudioTarget>
          <ChartCard.HeroNumber>{fmtUsd(model.marginalCostAtMax)}</ChartCard.HeroNumber>
          <p className="tabular text-sm text-muted-foreground">
            marginal cost / conversion
          </p>
        </ChartCard.HeaderMain>
        {(spec.legend === "full" || spec.legend === "inline") && (
          <ChartCard.Legend
            items={[
              { label: config.actual.label as string, color: actualColor },
              { label: config.fit.label as string, color: fitColor },
            ]}
          />
        )}
      </ChartCard.Header>

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

      <ChartCard.Body className={classNames?.body}>
        <MarketingSpendSaturationPlot
          data={data}
          config={config}
          heightClass={spec.height}
          valueFormatter={valueFormatter}
        />
      </ChartCard.Body>

      <StudioTarget fieldKey="attribution" variant="inline">
        <ChartCard.Footnote className={classNames?.footnote}>
          {attributionText}
        </ChartCard.Footnote>
      </StudioTarget>
    </>
  );
}

type PlotProps = {
  /** Raw weekly spend vs conversions — the √-fit + samples are derived inside. */
  data: SaturationPoint[];
  /** Series colors + labels (a shadcn ChartConfig keyed by `fit` / `actual`). */
  config: ChartConfig;
  heightClass?: string;
  /** Override conversion (Y) formatting. */
  valueFormatter?: (n: number) => string;
};

/**
 * Just the chart-drawing piece — the fitted response curve, the observed-week
 * scatter, the diminishing-returns marker, axes and tooltip, with no card
 * chrome, narrative or states. Derives the fit from the same raw data the widget
 * does, so you can drop it into your own layout.
 */
export function MarketingSpendSaturationPlot({
  data,
  config,
  heightClass = "h-64",
  valueFormatter,
}: PlotProps) {
  const { chartData, knee } = React.useMemo(() => buildSaturationModel(data), [data]);
  const fitColor = config.fit?.color ?? "var(--chart-1)";
  const actualColor = config.actual?.color ?? "var(--chart-3)";
  const fmt = valueFormatter ?? fmtConv;

  return (
    <ChartContainer config={config} className={cn(heightClass, "w-full")}>
      <ComposedChart data={chartData} margin={{ left: 0, right: 8, top: 8 }}>
        <CartesianGrid vertical={false} stroke="var(--chart-grid)" />
        <XAxis
          type="number"
          dataKey="spend"
          domain={["dataMin", "dataMax"]}
          tickFormatter={(v: number) => fmtUsd(v)}
          tickLine={false}
          axisLine={false}
          tickMargin={10}
          fontSize={10}
          minTickGap={32}
          className="smallcaps tabular"
        />
        <YAxis
          tickFormatter={(v: number) => fmt(v)}
          tickLine={false}
          axisLine={false}
          width={40}
          fontSize={10}
          className="smallcaps tabular"
        />
        <ChartTooltip
          cursor={{
            stroke: "var(--foreground)",
            strokeOpacity: 0.25,
            strokeDasharray: "2 4",
          }}
          content={
            <SaturationTooltip
              fmt={fmt}
              fitColor={fitColor}
              actualColor={actualColor}
            />
          }
        />
        <ReferenceLine
          x={knee}
          stroke="var(--foreground)"
          strokeOpacity={0.4}
          strokeDasharray="3 3"
          label={{
            value: "returns flatten",
            position: "top",
            fontSize: 9,
            fill: "var(--muted-foreground)",
          }}
        />
        <Line
          dataKey="fit"
          type="monotone"
          stroke={fitColor}
          strokeWidth={1.75}
          dot={false}
          isAnimationActive={false}
        />
        <Scatter dataKey="actual" fill={actualColor} isAnimationActive={false} />
      </ComposedChart>
    </ChartContainer>
  );
}