import { useState, useEffect, useRef } from "react";
function estimateTokens(charCount) {
  return Math.ceil(charCount / 4);
}
function formatTokenCount(count, includesEstimate) {
  if (count < 1e3) return "";
  const prefix = includesEstimate ? "~" : "";
  const kValue = (count / 1e3).toFixed(1);
  return `${prefix}${kValue}K tokens`;
}
function useTokenCounter({ agent, isActive }) {
  const [lastInputTokens, setLastInputTokens] = useState(0);
  const [cumulativeOutputTokens, setCumulativeOutputTokens] = useState(0);
  const [currentSegmentEstimate, setCurrentSegmentEstimate] = useState(0);
  const currentCharCountRef = useRef(0);
  useEffect(() => {
    if (!isActive) {
      setLastInputTokens(0);
      setCumulativeOutputTokens(0);
      setCurrentSegmentEstimate(0);
      currentCharCountRef.current = 0;
      return;
    }
    const controller = new AbortController();
    const { signal } = controller;
    currentCharCountRef.current = 0;
    setLastInputTokens(0);
    setCumulativeOutputTokens(0);
    setCurrentSegmentEstimate(0);
    agent.on(
      "llm:chunk",
      (payload) => {
        if (payload.chunkType === "text") {
          currentCharCountRef.current += payload.content.length;
          const estimate = estimateTokens(currentCharCountRef.current);
          if (estimate >= 1e3) {
            setCurrentSegmentEstimate(estimate);
          }
        }
      },
      { signal }
    );
    agent.on(
      "llm:response",
      (payload) => {
        const usage = payload.tokenUsage;
        const rawInputTokens = usage.inputTokens ?? 0;
        const cacheWriteTokens = usage.cacheWriteTokens ?? 0;
        const inputTokens = Math.max(0, rawInputTokens - cacheWriteTokens);
        if (inputTokens > 0) {
          setLastInputTokens(inputTokens);
        }
        const outputTokens = usage.outputTokens ?? 0;
        if (outputTokens > 0) {
          setCumulativeOutputTokens((prev) => prev + outputTokens);
        }
        currentCharCountRef.current = 0;
        setCurrentSegmentEstimate(0);
      },
      { signal }
    );
    return () => {
      controller.abort();
    };
  }, [agent, isActive]);
  const totalActualTokens = lastInputTokens + cumulativeOutputTokens;
  const displayCount = totalActualTokens + currentSegmentEstimate;
  const includesEstimate = currentSegmentEstimate > 0;
  return {
    totalActualTokens,
    currentSegmentEstimate,
    displayCount,
    includesEstimate,
    formatted: formatTokenCount(displayCount, includesEstimate)
  };
}
export {
  useTokenCounter
};
