frontend/src/components/chat/ChatBubble.tsx

import { useState } from "react";
import { ArrowRight, RefreshCcw, Copy, Check, Trash2, RotateCcw, ListFilter, ChevronLeft, ChevronRight } from "lucide-react";
import { Button } from "@/components/ui/button";
import { cn } from "@/lib/utils";
import { format } from "date-fns";
import { Avatar, AvatarFallback } from "../ui/avatar";
import { ChatMessage } from "./ChatMessage";
import { ThinkingAnimation } from "./ThinkingAnimation";
import { toast } from '../ui/sonner';
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogDescription, DialogFooter } from "@/components/ui/dialog";
import { AlertDialog, AlertDialogAction, AlertDialogCancel, AlertDialogContent, AlertDialogDescription, AlertDialogFooter, AlertDialogHeader, AlertDialogTitle } from "@/components/ui/alert-dialog";

interface MessageVariation {
  id: string;
  content: string;
  timestamp: Date;
}

interface Message {
  id: string;
  content: string;
  sender: "user" | "system";
  timestamp: Date;
  isLoading?: boolean;
  error?: boolean;
  result?: any;
  variations?: MessageVariation[];
  activeVariation?: string;
}

interface ChatBubbleProps {
  message: Message;
  onViewSearchResults?: (messageId: string) => void;
  onRetry?: (messageId: string) => void;
  onRegenerate?: (messageId: string) => void;
  onDelete?: (messageId: string) => void;
  onSelectVariation?: (messageId: string, variationId: string) => void;
}

export const ChatBubble = ({ 
  message, 
  onViewSearchResults, 
  onRetry, 
  onRegenerate,
  onDelete,
  onSelectVariation
}: ChatBubbleProps) => {
  const [copied, setCopied] = useState(false);
  const [deleteDialogOpen, setDeleteDialogOpen] = useState(false);
  const isWelcomeMessage = message.id === "welcomems"
  const isSystem = message.sender === "system";
  const showCopyButton = true;
  const hasVariations = isSystem && message.variations && message.variations.length > 0;
  
  // If the message has variations, display the active one or the first one
  const displayContent = isSystem && hasVariations && message.activeVariation
    ? message.variations.find(v => v.id === message.activeVariation)?.content || message.content
    : message.content;
  
  const copyToClipboard = () => {
    navigator.clipboard.writeText(displayContent);
    setCopied(true);
    toast.success('Copied to clipboard!');
    setTimeout(() => setCopied(false), 2000);
  };

  const handleDelete = () => {
    setDeleteDialogOpen(false);
    if (onDelete) {
      onDelete(message.id);
    }
  };

  const handleSelectVariation = (variationId: string) => {
    if (onSelectVariation) {
      onSelectVariation(message.id, variationId);
    }
  };

  // Function to get the current variation index and navigate through variations
  const navigateVariations = (direction: 'prev' | 'next') => {
    if (!hasVariations || !message.variations || message.variations.length <= 1) return;
    
    const currentIndex = message.activeVariation 
      ? message.variations.findIndex(v => v.id === message.activeVariation)
      : 0;
      
    let newIndex;
    if (direction === 'prev') {
      newIndex = (currentIndex - 1 + message.variations.length) % message.variations.length;
    } else {
      newIndex = (currentIndex + 1) % message.variations.length;
    }
    
    handleSelectVariation(message.variations[newIndex].id);
  };

  // Get current variation index for display
  const getCurrentVariationIndex = () => {
    if (!hasVariations || !message.variations) return 0;
    return message.activeVariation 
      ? message.variations.findIndex(v => v.id === message.activeVariation) + 1
      : 1;
  };

  return (
    <>
      <div
        className={cn(
          "group flex w-full animate-slide-in mb-4",
          isSystem ? "justify-start" : "justify-end"
        )}
      >
        <div className={cn(
          "flex gap-3 max-w-[80%]",
          isSystem ? "flex-row" : "flex-row-reverse"
        )}>
          <Avatar className={cn(
            "h-8 w-8 border",
            isSystem
              ? "bg-financial-accent dark:text-white shadow-lg"
              : "bg-muted"
          )}>
            <AvatarFallback className="text-xs font-semibold">
              {isSystem ? "AI" : "You"}
            </AvatarFallback>
          </Avatar>

          <div className="flex flex-col">
            <div className={cn(
              "rounded-2xl shadow-lg message-bubble backdrop-blur-sm",
              isSystem
                ? "bg-white/90 dark:bg-card/90 border border-border text-foreground message-bubble-ai"
                : "bg-financial-accent/30 border border-financial-accent/30 dark:text-white message-bubble-user",
              message.error && "border-destructive dark:border-red-500"
            )}>
              {message.isLoading ? (
                <ThinkingAnimation />
              ) : (
                <ChatMessage
                  content={displayContent}
                />
              )}
            </div>
            
            {/* Chat bubble footer */}
            <div className="flex flex-row chat-bubble-footer justify-between">
              {/* Time */}
              <div className={cn(
                "text-xs text-muted-foreground mt-1",
                isSystem ? "text-left" : "text-right"
              )}>
                {format(message.timestamp, "h:mm a")}
              </div>
              {/* Controls */}
              {!isWelcomeMessage && <div className="controls flex gap-1">
                {/* Variation Navigation */}
                {hasVariations && message.variations && message.variations.length > 1 && (
                  <div className="flex items-center mt-1 opacity-0 group-hover:opacity-100 transition-opacity bg-background/50 rounded-md border">
                    <Button
                      variant="ghost"
                      size="sm"
                      onClick={() => navigateVariations('prev')}
                      disabled={message.variations.length <= 1}
                      className="h-6 px-1"
                    >
                      <ChevronLeft className="h-3 w-3" />
                    </Button>
                    <span className="text-xs px-1">
                      {getCurrentVariationIndex()}/{message.variations.length}
                    </span>
                    <Button
                      variant="ghost"
                      size="sm"
                      onClick={() => navigateVariations('next')}
                      disabled={message.variations.length <= 1}
                      className="h-6 px-1"
                    >
                      <ChevronRight className="h-3 w-3" />
                    </Button>
                  </div>
                )}
              
                {/* Retry button for failed messages */}
                {message.error && onRetry && (
                  <div className="flex mt-1">
                    <Button
                      variant="secondary"
                      size="sm"
                      onClick={() => onRetry(message.id)}
                      className="text-xs flex items-center gap-1.5 text-muted-foreground hover:border border-financial-accent/30 bg-background/50 backdrop-blur-sm"
                    >
                      <RefreshCcw className="h-3 w-3" />
                      Retry
                    </Button>
                  </div>
                )}
                
                {/* Regenerate button for system messages */}
                {isSystem && !message.isLoading && onRegenerate && (
                  <TooltipProvider>
                    <Tooltip>
                      <TooltipTrigger asChild>
                        <div className="flex mt-1 opacity-0 group-hover:opacity-100 transition-opacity">
                          <Button
                            variant="link"
                            size="sm"
                            onClick={() => onRegenerate(message.id)}
                          >
                            <RotateCcw className="h-3 w-3" />
                          </Button>
                        </div>
                      </TooltipTrigger>
                      <TooltipContent>
                        <p>Generate variation</p>
                      </TooltipContent>
                    </Tooltip>
                  </TooltipProvider>
                )}
                
                {/* Delete button */}
                {onDelete && !message.isLoading && (
                  <TooltipProvider>
                    <Tooltip>
                      <TooltipTrigger asChild>
                        <div className="flex mt-1 opacity-0 group-hover:opacity-100">
                          <Button
                            variant="ghost"
                            size="sm"
                            onClick={() => setDeleteDialogOpen(true)}
                            className="text-xs flex items-center gap-1.5 text-muted-foreground hover:text-destructive hover:bg-background/50 backdrop-blur-sm"
                          >
                            <Trash2 className="h-3 w-3" />
                          </Button>
                        </div>
                      </TooltipTrigger>
                      <TooltipContent>
                        <p>Delete message</p>
                      </TooltipContent>
                    </Tooltip>
                  </TooltipProvider>
                )}
                
                {/* Copy */}
                {showCopyButton && !message.isLoading && (
                  <div className="relative top-1 right-1 opacity-0 group-hover:opacity-100 transition-opacity">
                    <Button variant="link" size="sm" onClick={copyToClipboard}>
                      {copied ? <Check className="h-4 w-4" /> : <Copy className="h-4 w-4" />}
                    </Button>
                  </div>
                )}
              </div>}
            </div>
          </div>
        </div>
      </div>
      
      {/* Delete confirmation dialog */}
      <AlertDialog open={deleteDialogOpen} onOpenChange={setDeleteDialogOpen}>
        <AlertDialogContent>
          <AlertDialogHeader>
            <AlertDialogTitle>Delete Message</AlertDialogTitle>
            <AlertDialogDescription>
              Are you sure you want to delete this message? This action cannot be undone.
            </AlertDialogDescription>
          </AlertDialogHeader>
          <AlertDialogFooter>
            <AlertDialogCancel>Cancel</AlertDialogCancel>
            <AlertDialogAction onClick={handleDelete} className="bg-destructive text-destructive-foreground hover:bg-destructive/90">
              Delete
            </AlertDialogAction>
          </AlertDialogFooter>
        </AlertDialogContent>
      </AlertDialog>
    </>
  );
};