import { type JSX, memo, useMemo } from "react";
import { cn, formatTokens } from "../../../lib/utils";
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "../../ui/tooltip";
import type { AnatomyRole, AnatomySegment } from "./types";

const ROLE_COLOR_CLASSES: Record<AnatomyRole, string> = {
  system: "bg-sky-500/70",
  user: "bg-emerald-500/70",
  assistant: "bg-violet-500/70",
  tool: "bg-amber-500/70",
  tools: "bg-slate-500/70",
};

const ROLE_FOCUS_RING: Record<AnatomyRole, string> = {
  system: "focus-visible:ring-sky-300",
  user: "focus-visible:ring-emerald-300",
  assistant: "focus-visible:ring-violet-300",
  tool: "focus-visible:ring-amber-300",
  tools: "focus-visible:ring-slate-300",
};

const MAX_VISIBLE_SEGMENTS = 12;
const MIN_SEGMENT_PERCENT = 1;

const TOOLTIP_PREVIEW_LIMIT = 80;
const LABEL_TRUNCATE_LIMIT = 24;

function truncateLabel(label: string): string {
  if (label.length <= LABEL_TRUNCATE_LIMIT) return label;
  return `${label.slice(0, LABEL_TRUNCATE_LIMIT - 1)}…`;
}

function truncatePreview(text: string): string {
  const singleLine = text.replace(/\s+/g, " ").trim();
  if (singleLine.length <= TOOLTIP_PREVIEW_LIMIT) return singleLine;
  return `${singleLine.slice(0, TOOLTIP_PREVIEW_LIMIT)}…`;
}

export type SegmentBarProps = {
  segments: ReadonlyArray<AnatomySegment>;
  totalTokens: number;
  onActivate?: (segment: AnatomySegment) => void;
};

/**
 * Render a horizontal stacked bar showing the relative size of each
 * request segment, with role-based color, per-segment labels, and
 * an overflow pill that aggregates any segments beyond the visible
 * cap. Click or keyboard activation (Enter / Space) calls `onActivate`.
 */
export const SegmentBar = memo(function SegmentBar({
  segments,
  totalTokens,
  onActivate,
}: SegmentBarProps): JSX.Element {
  const total = useMemo(() => {
    if (totalTokens > 0) return totalTokens;
    return segments.reduce((sum, s) => sum + s.size, 0);
  }, [segments, totalTokens]);

  const visibleSegments = segments.slice(0, MAX_VISIBLE_SEGMENTS);
  const overflowSegments = segments.slice(MAX_VISIBLE_SEGMENTS);
  const overflowSize = overflowSegments.reduce((sum, s) => sum + s.size, 0);
  const overflowCount = overflowSegments.length;
  const overflowStartIndex = MAX_VISIBLE_SEGMENTS;
  const overflowEndIndex = overflowStartIndex + overflowCount - 1;
  const hasOverflow = overflowCount > 0;

  const visibleTotal = useMemo(
    () => visibleSegments.reduce((sum, s) => sum + s.size, 0),
    [visibleSegments],
  );

  const ariaLabel = useMemo(
    () =>
      `Request anatomy: ~${formatTokens(total)} tokens across ${segments.length} segment${segments.length === 1 ? "" : "s"}`,
    [segments.length, total],
  );

  if (segments.length === 0 || total <= 0) {
    return <div role="img" aria-label={ariaLabel} className="h-6" />;
  }

  return (
    <TooltipProvider delayDuration={150}>
      <div role="img" aria-label={ariaLabel} className="flex flex-col gap-1.5">
        <div
          className="relative flex h-6 w-full overflow-hidden rounded border border-border/40 bg-muted/30"
          data-testid="anatomy-segment-bar"
        >
          {visibleSegments.map((segment, index) => {
            const rawPercent = total > 0 ? (segment.size / total) * 100 : 0;
            const percent = Math.max(MIN_SEGMENT_PERCENT, rawPercent);
            return (
              <Tooltip key={`${segment.path}-${index}`}>
                <TooltipTrigger asChild>
                  <button
                    type="button"
                    role="button"
                    tabIndex={0}
                    onClick={() => onActivate?.(segment)}
                    onKeyDown={(e) => {
                      if (e.key === "Enter" || e.key === " ") {
                        e.preventDefault();
                        onActivate?.(segment);
                      }
                    }}
                    data-anatomy-path={segment.path}
                    aria-label={`${segment.label}, ~${formatTokens(segment.size)} tokens`}
                    className={cn(
                      "h-full border-r border-background/80 last:border-r-0",
                      "opacity-90 hover:opacity-100 focus:opacity-100",
                      "focus:outline-none focus-visible:ring-2 focus-visible:ring-offset-1 focus-visible:ring-offset-background",
                      ROLE_COLOR_CLASSES[segment.role],
                      ROLE_FOCUS_RING[segment.role],
                    )}
                    style={{ width: `${percent}%` }}
                  />
                </TooltipTrigger>
                <TooltipContent side="bottom" className="max-w-sm text-xs p-2 space-y-0.5">
                  <div className="font-semibold">
                    {segment.label} · ~{formatTokens(segment.size)} tokens
                  </div>
                  <div className="text-muted-foreground">
                    {segment.characters.toLocaleString()} chars
                  </div>
                  {segment.text.length > 0 && (
                    <div className="text-muted-foreground/90 break-words whitespace-pre-wrap">
                      {truncatePreview(segment.text)}
                    </div>
                  )}
                </TooltipContent>
              </Tooltip>
            );
          })}
          {hasOverflow && (
            <Tooltip>
              <TooltipTrigger asChild>
                <div
                  role="img"
                  aria-label={`${overflowCount} additional segments`}
                  className="flex h-full items-center justify-center bg-muted-foreground/30 px-1.5 text-[10px] font-mono text-background"
                  style={{
                    width: `${Math.max(MIN_SEGMENT_PERCENT, (overflowSize / total) * 100)}%`,
                  }}
                >
                  … +{overflowCount}
                </div>
              </TooltipTrigger>
              <TooltipContent side="bottom" className="text-xs">
                {overflowCount} more segment{overflowCount === 1 ? "" : "s"} (indices{" "}
                {overflowStartIndex}–{overflowEndIndex})
              </TooltipContent>
            </Tooltip>
          )}
        </div>

        <div className="flex w-full gap-1 text-[10px] text-muted-foreground">
          {visibleSegments.map((segment, index) => {
            const rawPercent = total > 0 ? (segment.size / total) * 100 : 0;
            const percent = Math.max(MIN_SEGMENT_PERCENT, rawPercent);
            return (
              <div
                key={`label-${segment.path}-${index}`}
                className="flex flex-col gap-0.5 truncate"
                style={{ width: `${percent}%` }}
                title={`${segment.label} · ~${formatTokens(segment.size)} tokens`}
              >
                <span className="truncate font-mono text-foreground/80">
                  {truncateLabel(segment.label)}
                </span>
                <span className="truncate font-mono text-muted-foreground/70">
                  ~{formatTokens(segment.size)}
                </span>
              </div>
            );
          })}
          {hasOverflow && (
            <div
              className="flex flex-col gap-0.5 truncate text-muted-foreground"
              style={{ width: `${Math.max(MIN_SEGMENT_PERCENT, (overflowSize / total) * 100)}%` }}
            >
              <span className="truncate font-mono text-foreground/60">… +{overflowCount}</span>
              <span className="truncate font-mono text-muted-foreground/60">
                ~{formatTokens(overflowSize)}
              </span>
            </div>
          )}
        </div>

        {/* Spacer to maintain minimum bar height for tiny segments */}
        {visibleTotal < total * 0.1 && <div className="h-0" />}
      </div>
    </TooltipProvider>
  );
});
