frontend/src/components/TokenUsageDisplay.tsx

import { TokenUsage } from "@/types/chat";

interface Props {
  tokenUsage: TokenUsage;
  modelUsed: string;
  primaryModel?: string;
}

const MODEL_LABELS: Record<string, string> = {
  "llama-3.1-8b-instant":     "Llama 3.1 · 8B",
  "llama-3.3-70b-versatile":  "Llama 3.3 · 70B",
  "mixtral-8x7b-32768":       "Mixtral · 8×7B",
  "gemma2-9b-it":             "Gemma 2 · 9B",
};

function shortLabel(model: string): string {
  return MODEL_LABELS[model] ?? model.split("-").slice(0, 3).join("-");
}

export default function TokenUsageDisplay({ tokenUsage, modelUsed, primaryModel }: Props) {
  const isFallback = primaryModel && modelUsed !== primaryModel;
  const { prompt_tokens, completion_tokens, total_tokens } = tokenUsage;

  // Ratio bar: completion vs prompt (capped at 100%)
  const ratio = total_tokens > 0 ? Math.min((completion_tokens / total_tokens) * 100, 100) : 0;

  return (
    <div className="flex items-center gap-3 mt-2.5 flex-wrap">
      {/* Model badge */}
      <div className={`flex items-center gap-1.5 rounded-full border px-2.5 py-0.5 ${
        isFallback
          ? "border-amber-200 bg-amber-50 text-amber-600"
          : "border-[#E8EDF2] bg-[#F8FAFC] text-[#64748B]"
      }`}>
        {isFallback ? (
          <svg width="9" height="9" viewBox="0 0 10 10" fill="none" className="flex-shrink-0">
            <path d="M5 1v4l2.5 1.5" stroke="currentColor" strokeWidth="1.5" strokeLinecap="round"/>
            <circle cx="5" cy="5" r="4" stroke="currentColor" strokeWidth="1.2"/>
          </svg>
        ) : (
          <svg width="9" height="9" viewBox="0 0 10 10" fill="currentColor" className="flex-shrink-0 opacity-60">
            <circle cx="5" cy="5" r="3"/>
          </svg>
        )}
        <span className="text-[10px] font-semibold leading-none">
          {shortLabel(modelUsed)}
        </span>
        {isFallback && (
          <span className="text-[9px] font-medium uppercase tracking-wider opacity-70 leading-none">
            fallback
          </span>
        )}
      </div>

      {/* Token breakdown */}
      <div className="flex items-center gap-1.5">
        {/* Mini bar */}
        <div className="w-12 h-1.5 rounded-full bg-[#E8EDF2] overflow-hidden">
          <div
            className="h-full rounded-full bg-gradient-to-r from-[#47C3A6] to-[#14B7CC]"
            style={{ width: `${ratio}%` }}
          />
        </div>

        <span className="text-[10px] text-[#94A3B8] font-mono tabular-nums">
          <span title="Prompt tokens" className="text-[#94A3B8]">{prompt_tokens.toLocaleString()}</span>
          <span className="text-[#CBD5E1] mx-0.5">+</span>
          <span title="Completion tokens" className="text-[#47C3A6]">{completion_tokens.toLocaleString()}</span>
          <span className="text-[#CBD5E1] mx-1">·</span>
          <span title="Total tokens" className="font-semibold text-[#64748B]">{total_tokens.toLocaleString()}</span>
          <span className="text-[#CBD5E1] ml-0.5">tok</span>
        </span>
      </div>
    </div>
  );
}