"use client";

import { useCallback, useEffect, useLayoutEffect, useRef, useState } from "react";

import {
  Model,
  Gpu,
  PRECISIONS,
  PRECISION_ACCENT,
  overheadGB,
  round1,
  weightsGB,
} from "@/lib/data";

/** Width before the container has been measured. */
const DEFAULT_W = 1180;
const BAR_H = 55;
const R = 5.5;
/** Pointer band around each tick, in viewBox units. */
const HIT_W = 44;
/** MethodReceipt's fixed width. */
const CARD_W = 466;

/**
 * Where the hover card's left edge goes, in px from the chart's left edge, so
 * that it reads as hanging from the tick at `anchorX` without running off a
 * chart `width` px wide. The chart is as narrow as the docs iframe (~730px),
 * so a card can fail to fit on either side of a tick in the middle.
 */
function cardLeft(anchorX: number, width: number, cardW: number): number {
  // Hang from the tick: rightward if the card fits, else leftward.
  const hung = anchorX + cardW <= width ? anchorX : anchorX - cardW;
  // Neither side fits, so slide it into frame. When the chart is narrower than
  // the card, the left edge wins: overflow to the right can still be scrolled.
  return Math.max(0, Math.min(hung, width - cardW));
}

interface Props {
  /** null until a model resolves — the bar renders empty, but to scale */
  model: Model | null;
  gpu: Gpu;
  contextTokens: number;
  selectedBits: number;
  /** fires with the precision under the cursor, or null on leave */
  onHoverBits?: (bits: number | null) => void;
  /** which tick is under the cursor, so the bar can anchor the card to it */
  hoverBits?: number | null;
  /** rendered under the hovered tick, positioned by the bar */
  hoverCard?: React.ReactNode;
}

interface Tick {
  bits: number;
  total: number;
  x: number;
  /** which label row this tick's caption sits on, 0 = closest to the bar */
  row: number;
}

/** Round a scale ceiling up to something a human would pick. */
function niceCeil(v: number): number {
  const step = v > 400 ? 100 : v > 200 ? 50 : v > 60 ? 10 : 5;
  return Math.ceil(v / step) * step;
}

/**
 * The chart draws in real pixels: the viewBox is as wide as the box it sits
 * in. A fixed viewBox would scale the 11px captions down with the page, and in
 * the docs iframe they shrink to about 6px.
 */
function useWidth() {
  const ref = useRef<HTMLDivElement>(null);
  const [width, setWidth] = useState(DEFAULT_W);
  // Layout effect, so the first paint is already at the measured width.
  useLayoutEffect(() => {
    const el = ref.current;
    if (!el) return;
    setWidth(Math.max(1, Math.round(el.clientWidth)));
    const ro = new ResizeObserver(([entry]) => {
      setWidth(Math.max(1, Math.round(entry.contentRect.width)));
    });
    ro.observe(el);
    return () => ro.disconnect();
  }, []);
  return [ref, width] as const;
}

export default function MemoryBar(props: Props) {
  const [ref, W] = useWidth();
  return (
    <div ref={ref}>
      {props.model ? <FilledBar {...props} model={props.model} W={W} /> : <EmptyBar W={W} />}
    </div>
  );
}

/**
 * The card, waiting. Just the empty track, holding the same space the filled
 * bar will take, so nothing moves when a model arrives.
 */
function EmptyBar({ W }: { W: number }) {
  const barTop = 84;
  const barBottom = barTop + BAR_H;

  return (
    <svg viewBox={`0 0 ${W} ${barBottom + 66}`} className="h-auto w-full" role="img"
      aria-label="No model chosen yet">
      <rect x="0.5" y={barTop + 0.5} width={W - 1} height={BAR_H} rx={R} fill="#ffffff"
        stroke="#ebebeb" />
    </svg>
  );
}

function FilledBar({ model, gpu, contextTokens, selectedBits, onHoverBits, hoverBits, hoverCard, W }: Props & { model: Model; W: number }) {
  const overhead = overheadGB(model, contextTokens);
  const vram = gpu.vram;

  // The solid bar is the precision you picked. fp16 trails behind it as a
  // ghost, so you can see what quantizing bought you.
  const selWeights = weightsGB(model, selectedBits);
  const selTotal = selWeights + overhead;
  const ghostTotal = weightsGB(model, 16) + overhead;

  // The scale follows the model, not the card: a 2B model on an 80GB card would
  // otherwise crowd all four precisions into a few pixels. It spans the widest
  // precision on offer, so every tick stays readable. fp16 and — when the model
  // fits comfortably — the capacity mark are free to run off the end.
  const widestTick = Math.max(...PRECISIONS.map((b) => weightsGB(model, b) + overhead));
  const scaleMax = niceCeil(Math.max(selTotal, widestTick) * 1.02);
  const x = (gb: number) => (gb / scaleMax) * W;

  const overflows = selTotal > vram;
  // When it overflows, the card is full: weights fill whatever the reservation leaves.
  const weightsThatFit = Math.max(0, vram - overhead);
  const kvStart = overflows ? x(weightsThatFit) : x(selWeights);
  const kvEnd = overflows ? x(vram) : x(selTotal);
  const barEnd = x(selTotal);
  const ghostEnd = x(ghostTotal);
  const ghostClipped = ghostEnd > W;
  const ghostRight = Math.min(ghostEnd, W - 0.5);
  // The track's corners are rounded, so a clipped ghost must stop before the
  // curve begins or its horizontals cut across the corner.
  const ghostRunEnd = ghostClipped ? W - R - 0.5 : ghostRight;
  // The card only needs showing while it is in play.
  const capX = x(vram);
  const capOff = capX > W;
  // The callout needs ~215px; right-align it on the block when it would run off.
  const kvLabelFlips = kvStart + 215 > W;
  // Two totals this close would overprint.
  const showGhostLabel = ghostRunEnd - barEnd >= 64;

  // Precision labels are ~56px wide; anything closer than that overprints. Each
  // tick takes the first row with room, so a small model — whose four
  // precisions can land within a few pixels of each other — fans out upward
  // instead of stacking into one illegible blur. The lines stay at their true
  // x: only the captions move.
  const MIN_GAP = 56;
  const ROW_H = 42;
  const lastX: number[] = [];
  const ticks: Tick[] = PRECISIONS.map((bits) => {
    const total = weightsGB(model, bits) + overhead;
    return { bits, total, x: x(total), row: 0 };
  })
    .filter((t) => t.total <= scaleMax)
    .sort((a, b) => a.x - b.x)
    .map((t) => {
      let row = 0;
      while (lastX[row] !== undefined && t.x - lastX[row] < MIN_GAP) row++;
      lastX[row] = t.x;
      return { ...t, row };
    });

  // Row 0 is the topmost caption and has the longest drop line; each extra row
  // sits 42px lower, so more rows push the bar down.
  const rows = Math.max(1, lastX.length);
  const barTop = Math.max(84, 40 + (rows - 1) * ROW_H + 18);
  const barBottom = barTop + BAR_H;
  const height = barBottom + 66;

  // The card hangs below the chart, outside this container's box, so pointing
  // at it would otherwise read as leaving. Closing waits a beat, and anything
  // the pointer lands on in the meantime cancels it.
  const closeTimer = useRef<ReturnType<typeof setTimeout> | undefined>(undefined);
  const cancelClose = useCallback(() => clearTimeout(closeTimer.current), []);
  const scheduleClose = useCallback(() => {
    clearTimeout(closeTimer.current);
    closeTimer.current = setTimeout(() => onHoverBits?.(null), 160);
  }, [onHoverBits]);
  useEffect(() => () => clearTimeout(closeTimer.current), []);

  const anchor = ticks.find((t) => t.bits === hoverBits);

  return (
    <div className="relative">
    <svg viewBox={`0 0 ${W} ${height}`} className="h-auto w-full" role="img"
      aria-label={`${model.label} at ${selectedBits}-bit needs ${round1(selTotal)} GB of the ${vram} GB on ${gpu.label}`}>
      {/* precision labels + drop lines above the bar */}
      {ticks.map((t) => {
        const accent = PRECISION_ACCENT[t.bits];
        const isSel = t.bits === selectedBits;
        const nameY = 11 + t.row * ROW_H;
        const valY = 29 + t.row * ROW_H;
        const lineTop = 40 + t.row * ROW_H;
        return (
          <g key={t.bits}>
            <text x={t.x} y={nameY} textAnchor="middle" fontFamily="var(--font-geist-mono)"
              fontSize="11" letterSpacing="0.79" fill={accent}>
              {t.bits}-BIT
            </text>
            <text x={t.x} y={valY} textAnchor="middle" fontFamily="var(--font-geist-mono)"
              fontSize="13" fontWeight={isSel ? 500 : 400} fill={isSel ? accent : "#171717"}>
              {round1(t.total)}
            </text>
            <line x1={t.x} y1={lineTop} x2={t.x} y2={barTop} stroke={accent}
              strokeWidth={isSel ? 2.5 : 1.75} strokeLinecap="round"
              strokeDasharray={isSel ? "0.01 4" : "0.01 3.5"}
              opacity={isSel ? 1 : 0.9} />
          </g>
        );
      })}

      {/* KV reservation callout — under the bar, anchored on the block it describes */}
      <text
        x={kvLabelFlips ? kvEnd : kvStart}
        y={barBottom + 52}
        textAnchor={kvLabelFlips ? "end" : "start"}
        fontFamily="var(--font-geist-mono)"
        fontSize="11"
        letterSpacing="0.79"
        fill="#8f8f8f"
      >
        <tspan fontSize="12" letterSpacing="0" fill="#666666">{round1(overhead)}</tspan>
        {"  KV CACHE + ACTIVATIONS"}
      </text>

      {/* track */}
      <rect x="0.5" y={barTop + 0.5} width={W - 1} height={BAR_H} rx={R} fill="#ffffff"
        stroke="#ebebeb" />

      {/* fp16, unquantized — the ghost the solid bar is measured against. It is
          left open on the right when it runs past the edge of the scale. */}
      <path
        d={
          ghostClipped
            ? `M${ghostRunEnd} ${barTop + 0.5} H${R} A${R} ${R} 0 0 0 0.5 ${barTop + R} V${barBottom - R} A${R} ${R} 0 0 0 ${R} ${barBottom + 0.5} H${ghostRunEnd}`
            : `M${ghostRight} ${barTop + 0.5} H${R} A${R} ${R} 0 0 0 0.5 ${barTop + R} V${barBottom - R} A${R} ${R} 0 0 0 ${R} ${barBottom + 0.5} H${ghostRight} Z`
        }
        fill="none"
        stroke="#d4d4d4"
        strokeWidth="1"
        strokeDasharray="2 3"
      />
      {showGhostLabel && (
        <text x={ghostRunEnd - 6} y={barTop - 8} textAnchor="end"
          fontFamily="var(--font-geist-mono)" fontSize="11" letterSpacing="0.79" fill="#a8a8a8">
          {round1(ghostTotal)} AT FP16{ghostClipped ? " ›" : ""}
        </text>
      )}

      {/* weights that fit — rounded on the left only */}
      <path
        d={`M${R} ${barTop} H${kvStart} V${barBottom} H${R} A${R} ${R} 0 0 1 0 ${barBottom - R} V${barTop + R} A${R} ${R} 0 0 1 ${R} ${barTop} Z`}
        fill="#171717"
      />
      {/* the fixed reservation */}
      <rect x={kvStart} y={barTop} width={Math.max(0, kvEnd - kvStart)} height={BAR_H} fill="#8f8f8f" />
      {/* what will not load */}
      {overflows && (
        <rect x={x(vram)} y={barTop} width={barEnd - x(vram)} height={BAR_H} fill="#ebebeb" />
      )}

      {/* drop lines continued through the bar */}
      {ticks.map((t) => {
        const accent = PRECISION_ACCENT[t.bits];
        const isSel = t.bits === selectedBits;
        return (
          <line key={t.bits} x1={t.x} y1={barTop + 1} x2={t.x} y2={barBottom - 1} stroke={accent}
            strokeWidth={isSel ? 2.5 : 1.75} strokeLinecap="round"
            strokeDasharray={isSel ? "0.01 4" : "0.01 3.5"}
            opacity={isSel ? 1 : 0.9} />
        );
      })}

      {/* hit targets — a 2px line is not a pointer target, so give each tick a
          band running the full height of its label and line */}
      {ticks.map((t) => (
        <rect
          key={`hit-${t.bits}`}
          x={t.x - HIT_W / 2}
          y={0}
          width={HIT_W}
          height={barBottom}
          fill="transparent"
          className="cursor-pointer outline-none"
          onMouseEnter={() => {
            cancelClose();
            onHoverBits?.(t.bits);
          }}
          onMouseLeave={scheduleClose}
          onFocus={() => {
            cancelClose();
            onHoverBits?.(t.bits);
          }}
          tabIndex={0}
          role="button"
          aria-label={`${t.bits}-bit, ${round1(t.total)} GB`}
        />
      ))}

      {/* in-bar figures */}
      {kvStart >= 52 && (
        <text x="16" y={barTop + 32} fontFamily="var(--font-geist-mono)" fontSize="13" fill="#fafafa">
          {round1(overflows ? weightsThatFit : selWeights)}
        </text>
      )}
      {overflows && barEnd - x(vram) >= 52 && (
        <text x={(x(vram) + barEnd) / 2} y={barTop + 32} textAnchor="middle"
          fontFamily="var(--font-geist-mono)" fontSize="13" fill="#171717">
          {round1(selTotal - vram)}
        </text>
      )}

      {/* capacity line — shown only while the card is in frame; past the edge
          it would just be a number with nothing to point at */}
      {!capOff && (
        <>
          <line x1={capX} y1={barBottom} x2={capX} y2={barBottom + 12} stroke="#171717" strokeWidth="1.5" />
          <text x={capX} y={barBottom + 30} textAnchor="middle" fontFamily="var(--font-geist-mono)"
            fontSize="13" fill="#171717">
            {vram}
          </text>
        </>
      )}

    </svg>

      {anchor && hoverCard && (
        <div
          className="absolute z-10"
          onMouseEnter={cancelClose}
          onMouseLeave={scheduleClose}
          style={{
            left: cardLeft(anchor.x, W, CARD_W),
            // Below the whole chart, so it never covers the KV cache caption.
            top: "calc(100% + 6px)",
          }}
        >
          {hoverCard}
        </div>
      )}
    </div>
  );
}