frontend/src/components/chat/ChatMessage.tsx

// src/components/chat/ChatMessage.tsx
import React, { useState, useMemo } from 'react'
import ReactMarkdown from 'react-markdown'
import remarkGfm from 'remark-gfm'
import rehypeRaw from 'rehype-raw'
import { cn } from '@/lib/utils'
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from '../ui/collapsible'
import { ChevronDown, Brain } from 'lucide-react'

interface ChatMessageProps {
  content: string
  className?: string
}

export const ChatMessage: React.FC<ChatMessageProps> = ({
  content,
  className,
}) => {
  // Extract thinking content and actual response
  const { processedContent, thinkingBlocks } = useMemo(() => {
    const blocks: { id: number; content: string }[] = [];
    let thinkBlockCounter = 0;

    // Extract thinking content between <think> tags
    const contentWithoutThinking = content.replace(
      /<think>([\s\S]*?)<\/think>/g,
      (_, thinkContent) => {
        blocks.push({ 
          id: thinkBlockCounter++, 
          content: thinkContent.trim() 
        });
        return ''; // Remove thinking content from the main message
      }
    );

    // Continue processing source tags
    const processedText = contentWithoutThinking.replace(
      /<source\s+path=["'](.+?)["']\s*\/>/g,
      (_match, path) => {
        const filename = path
          .split('/')
          .pop()!
          .replace(/\.[^/.]+$/, '')
        return `<a href="${path}" target="_blank" class="inline-flex items-center text-xs font-medium mx-0.5 rounded-sm px-1 bg-financial-accent/10 text-financial-accent border border-financial-accent/20 hover:bg-financial-accent/20 transition-colors">
            ${filename}
            <svg xmlns="http://www.w3.org/2000/svg" class="h-3 w-3 ml-1" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
              <path stroke-linecap="round" stroke-linejoin="round" d="M18 13v6a2 2 0 0 1-2 2H6a2 2 0 0 1-2-2V8a2 2 0 0 1 2-2h6M15 3h6m0 0v6m0-6L10 14"/>
            </svg>
          </a>`;
      }
    );

    return { 
      processedContent: processedText.trim(), 
      thinkingBlocks: blocks 
    };
  }, [content]);

  return (
    <div
      className={cn(
        'group relative w-full rounded-md p-2 hover:bg-muted/30 transition-colors',
        className
      )}
    >
      {/* First render the thinking blocks if any */}
      {thinkingBlocks.length > 0 && (
        <Collapsible 
          className="think-collapsible my-3 rounded-lg bg-financial-accent/5 mb-4"
          defaultOpen={false}
        >
          <CollapsibleTrigger className="flex items-center gap-2 w-full p-2 text-left hover:bg-financial-accent/10 rounded-t-lg">
            <div className="flex items-center gap-2 w-full">
              <div className="thinking-brain-small relative">
                <Brain className="h-4 w-4 text-financial-accent" />
              </div>
              <span className="text-xs font-medium text-financial-accent">Thoughts</span>
              <ChevronDown className="h-4 w-4 text-financial-accent/70 transition-transform duration-200 ml-auto" />
            </div>
          </CollapsibleTrigger>
          
          <CollapsibleContent>
            <div className="think-block p-3 text-sm text-muted-foreground bg-financial-accent/5">
              {thinkingBlocks.map((block, index) => (
                <ReactMarkdown 
                  key={`thinking-${block.id}`}
                  remarkPlugins={[remarkGfm]} 
                  rehypePlugins={[rehypeRaw]}
                >
                  {block.content}
                </ReactMarkdown>
              ))}
            </div>
          </CollapsibleContent>
        </Collapsible>
      )}

      {/* Then render the actual response content */}
      {processedContent && (
        <ReactMarkdown
          remarkPlugins={[remarkGfm]}
          rehypePlugins={[rehypeRaw]}
          components={{
            a: ({ href, children, node, ...props }) =>
              href && href.endsWith('.md') ? (
                <a
                  href={href}
                  target="_blank"
                  rel="noopener noreferrer"
                  {...props}
                >
                  {children}
                </a>
              ) : (
                <a href={href} {...props}>
                  {children}
                </a>
              ),
          }}
        >
          {processedContent}
        </ReactMarkdown>
      )}
    </div>
  )
}