export interface HighlightSegment {
  text: string;
  className: string;
}

export function highlightTrace(text: string): HighlightSegment[] {
  if (!text) return [];

  const segments: HighlightSegment[] = [];
  // Strip <think> and </think> tags for display
  const cleaned = text.replace(/<\/?think>/g, "");
  const lines = cleaned.split("\n");

  for (let i = 0; i < lines.length; i++) {
    const line = lines[i];
    const lo = line.toLowerCase().trim();

    let className = "text-gray-400";

    if (lo.startsWith("wait") || lo.startsWith("hmm") || lo.startsWith("but wait")) {
      className = "text-yellow-400";
    } else if (lo.startsWith("let me try") || lo.startsWith("let me reconsider") || lo.startsWith("let me think")) {
      className = "text-cyan-400";
    } else if (lo.startsWith("so the answer") || lo.startsWith("therefore") || lo.startsWith("the final") || lo.startsWith("i should")) {
      className = "text-green-400 font-bold";
    } else if (lo.startsWith("i give up") || lo.startsWith("i can't") || lo.startsWith("i'm stuck")) {
      className = "text-red-400 font-bold";
    } else if (line.includes("=") && /[+\-*/]/.test(line)) {
      className = "text-gray-200";
    }

    segments.push({ text: line, className });
    if (i < lines.length - 1) {
      segments.push({ text: "\n", className: "" });
    }
  }

  return segments;
}