import * as React from 'react';
import { Ansi, Box, Text, useAnimationFrame } from '@anthropic/ink';
import { segmentTextByHighlights, type TextHighlight } from '../../utils/textHighlighting.js';
import { ShimmerChar } from '../Spinner/ShimmerChar.js';

type Props = {
  text: string;
  highlights: TextHighlight[];
};

type LinePart = {
  text: string;
  highlight: TextHighlight | undefined;
  start: number;
};

export function HighlightedInput({ text, highlights }: Props): React.ReactNode {
  // The shimmer animation (below) re-renders this component at 20fps while the
  // ultrathink keyword is present. text/highlights are referentially stable
  // across animation ticks (parent doesn't re-render), so memoize everything
  // that derives from them: segmentTextByHighlights alone is ~85µs/call
  // (tokenize + sort + O(n²) overlap), which adds up fast at 20fps.
  const { lines, hasShimmer, sweepStart, cycleLength } = React.useMemo(() => {
    const segments = segmentTextByHighlights(text, highlights);

    // Split segments by newlines into per-line groups. Ink's row-direction Box
    // indents continuation lines of a multi-line child to that child's X offset.
    // By splitting at newlines, each line renders as its own row, avoiding the
    // incorrect indentation when highlighted text is followed by wrapped content.
    const lines: LinePart[][] = [[]];
    let pos = 0;
    for (const segment of segments) {
      const parts = segment.text.split('\n');
      for (let i = 0; i < parts.length; i++) {
        if (i > 0) {
          lines.push([]);
          pos += 1;
        }
        const part = parts[i]!;
        if (part.length > 0) {
          lines[lines.length - 1]!.push({
            text: part,
            highlight: segment.highlight,
            start: pos,
          });
        }
        pos += part.length;
      }
    }

    // Scope the sweep to shimmer-highlighted ranges so cycle time doesn't grow
    // with input length. Padding creates an offscreen pause between sweeps.
    const hasShimmer = highlights.some(h => h.shimmerColor);
    let sweepStart = 0;
    let cycleLength = 1;
    if (hasShimmer) {
      const padding = 10;
      let lo = Infinity;
      let hi = -Infinity;
      for (const h of highlights) {
        if (h.shimmerColor) {
          lo = Math.min(lo, h.start);
          hi = Math.max(hi, h.end);
        }
      }
      sweepStart = lo - padding;
      cycleLength = hi - lo + padding * 2;
    }

    return { lines, hasShimmer, sweepStart, cycleLength };
  }, [text, highlights]);

  const [ref, time] = useAnimationFrame(hasShimmer ? 50 : null);
  const glimmerIndex = hasShimmer ? sweepStart + (Math.floor(time / 50) % cycleLength) : -100;

  return (
    <Box ref={ref} flexDirection="column">
      {lines.map((lineParts, lineIndex) => (
        <Box key={lineIndex}>
          {lineParts.length === 0 ? (
            <Text> </Text>
          ) : (
            lineParts.map((part, partIndex) => {
              if (part.highlight?.shimmerColor && part.highlight.color) {
                return (
                  <Text key={partIndex}>
                    {part.text.split('').map((char, charIndex) => (
                      <ShimmerChar
                        key={charIndex}
                        char={char}
                        index={part.start + charIndex}
                        glimmerIndex={glimmerIndex}
                        messageColor={part.highlight!.color!}
                        shimmerColor={part.highlight!.shimmerColor!}
                      />
                    ))}
                  </Text>
                );
              }
              return (
                <Text
                  key={partIndex}
                  color={part.highlight?.color}
                  dimColor={part.highlight?.dimColor}
                  inverse={part.highlight?.inverse}
                >
                  <Ansi>{part.text}</Ansi>
                </Text>
              );
            })
          )}
        </Box>
      ))}
    </Box>
  );
}