import { Arrow, FigBox, FILL, LoopFigure, type LoopStop } from "./LoopFigure";

const CX = 300;

const STOPS: LoopStop[] = [
  {
    x: CX,
    y: 448,
    ids: ["scale"],
    inspect: "scale",
    title: "Prompt tokens",
    note: "A single token stream: instructions, context, and the answer so far. Scale is tokens × parameters × compute.",
  },
  {
    x: CX,
    y: 368,
    ids: ["decoder"],
    inspect: "decoder",
    title: "Embedding + positions",
    note: "Tokens become vectors. Causal masking will keep each position from seeing the future.",
  },
  {
    x: CX,
    y: 278,
    ids: ["decoder"],
    inspect: "decoder",
    title: "Masked self-attention",
    note: "Decoder-only: no encoder. Every layer is causal multi-head attention plus a feed-forward net.",
  },
  {
    x: CX,
    y: 198,
    ids: ["decoder"],
    inspect: "decoder",
    title: "Feed-forward × N",
    note: "Stacked N times. GPT-style models drop the encoder entirely and generate left to right.",
  },
  {
    x: CX,
    y: 108,
    ids: ["decoder"],
    inspect: "decoder",
    title: "Linear + softmax",
    note: "Next-token distribution. The sampled token is appended and the loop continues.",
  },
  {
    x: 560,
    y: 108,
    ids: ["align"],
    inspect: "align",
    title: "Post-training",
    note: "Instruction data and preference methods (RLHF, DPO) turn a next-token model into something you can talk to.",
  },
  {
    x: 560,
    y: 448,
    ids: ["align", "scale"],
    inspect: "align",
    title: "Append next token",
    note: "The predicted token joins the prompt. The decoder runs again.",
  },
];

export function LlmLoop({ onSelect }: { onSelect: (id: string) => void }) {
  return (
    <LoopFigure
      stops={STOPS}
      loopMs={15000}
      viewBox="0 0 700 500"
      ariaLabel="Decoder-only large language model architecture"
      onSelect={onSelect}
    >
      {({ on, markerId }) => (
        <g>
          <rect width="700" height="500" fill="#f7f4ec" />
          <text x="300" y="28" textAnchor="middle" fontSize="16" fontWeight="700" fill="#222">
            Large Language Model
          </text>
          <text x="300" y="48" textAnchor="middle" fontSize="12" fill="#555">
            Decoder-only Transformer
          </text>

          <FigBox
            x={450}
            y={86}
            w={220}
            h={44}
            fill={FILL.pink}
            lit={on("align")}
            label="Chat / tools / RLHF"
            onClick={() => onSelect("align")}
          />

          <rect
            x="170"
            y="156"
            width="260"
            height="160"
            rx="16"
            fill="none"
            stroke="#222"
            strokeWidth="1.4"
          />
          <text x="440" y="176" fontSize="13" fill="#222">
            Nx
          </text>

          <FigBox
            x={190}
            y={86}
            w={220}
            h={44}
            fill={FILL.green}
            lit={on("decoder")}
            label="Linear + Softmax"
            onClick={() => onSelect("decoder")}
          />
          <Arrow x1={CX} y1={156} x2={CX} y2={130} active={on("decoder")} markerId={markerId} />

          <FigBox
            x={190}
            y={176}
            w={220}
            h={44}
            fill={FILL.cyan}
            lit={on("decoder")}
            label="Feed Forward"
            onClick={() => onSelect("decoder")}
          />
          <Arrow x1={CX} y1={238} x2={CX} y2={220} active={on("decoder")} markerId={markerId} />
          <FigBox
            x={190}
            y={256}
            w={220}
            h={48}
            fill={FILL.orange}
            lit={on("decoder")}
            label="Masked Multi-Head"
            onClick={() => onSelect("decoder")}
          />

          <Arrow x1={CX} y1={338} x2={CX} y2={304} active={on("decoder")} markerId={markerId} />
          <FigBox
            x={190}
            y={346}
            w={220}
            h={44}
            fill={FILL.yellow}
            lit={on("decoder")}
            label="Embed + positions"
            onClick={() => onSelect("decoder")}
          />
          <Arrow x1={CX} y1={418} x2={CX} y2={390} active={on("scale") || on("decoder")} markerId={markerId} />
          <FigBox
            x={190}
            y={426}
            w={220}
            h={44}
            fill={FILL.purple}
            lit={on("scale")}
            label="Prompt tokens"
            onClick={() => onSelect("scale")}
          />

          <text x="80" y="230" textAnchor="middle" fontSize="11" fill="#555">
            params
          </text>
          <text x="80" y="248" textAnchor="middle" fontSize="11" fill="#555">
            × data
          </text>
          <text x="80" y="266" textAnchor="middle" fontSize="11" fill="#555">
            × compute
          </text>
          <path
            d="M125 248 H 170"
            fill="none"
            stroke={on("scale") ? "#111" : "#999"}
            strokeDasharray="5 4"
          />

          <path
            d="M410 108 H 450"
            fill="none"
            stroke={on("align") || on("decoder") ? "#111" : "#888"}
            strokeWidth={on("align") ? 2.2 : 1.2}
            markerEnd={`url(#${markerId})`}
          />
          <path
            d="M560 130 V 448 H 410"
            fill="none"
            stroke="#c45c5c"
            strokeWidth="1.2"
            strokeDasharray="5 6"
            opacity="0.7"
            markerEnd={`url(#${markerId})`}
          />
          <text x="620" y="290" fontSize="11" fill="#666">
            append
          </text>
          <text x="620" y="306" fontSize="11" fill="#666">
            token
          </text>
        </g>
      )}
    </LoopFigure>
  );
}