src/App.tsx

import { useState, useRef, useEffect, useMemo } from 'react';
import { ArrowUp, Bot, Trash2, Copy, Check, ImageUp, X, Brain, Pencil, Globe, Download, Wand2 } from 'lucide-react';
import { Sun, Moon } from 'lucide-react';
import { 
  GROQ_API_KEYS,
  GROQ_API_URL,
  GROQ_MODEL_REASONER,
  POLLINATIONS_API_URL,
  POLLINATIONS_MODEL_TEXT,
  POLLINATIONS_MODEL_VISION,
  POLLINATIONS_MODEL_FALLBACK,
  POLLINATIONS_MODEL_SEARCH,
  TACTIQ_API_URL
} from './config/api';
import ReactMarkdown, { type Components } from 'react-markdown';
import remarkGfm from 'remark-gfm';
import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
import { materialDark } from 'react-syntax-highlighter/dist/esm/styles/prism';
import 'katex/dist/katex.min.css';
import { InlineMath, BlockMath } from 'react-katex';
import type { Message, EditingMessage } from './types/msg';
import { SYSTEM_PROMPT, STORAGE_KEY, MAX_RETRIES, RETRY_DELAY, REQUEST_TIMEOUT } from './config/const';


const CodeBlock = ({ children, inline, className }: { children: string | string[]; inline?: boolean; className?: string }) => {
  const codeContent = Array.isArray(children) ? children.join('') : children?.toString() || '';
  
  if (inline || !className) {
    return <code className="px-1 py-0.5 rounded text-sm font-mono custom-scrollbar">{codeContent}</code>;
  }

  const [copied, setCopied] = useState(false);
  const language = className?.replace(/language-/, '') || 'plaintext';

  const handleCopy = async () => {
    await navigator.clipboard.writeText(codeContent);
    setCopied(true);
    setTimeout(() => setCopied(false), 2000);
  };

  const customStyle = {
    ...materialDark,
    'pre[class*="language-"]': {
      ...materialDark['pre[class*="language-"]'],
      background: 'transparent'
    }
  };

  return (
    <div className="relative group bg-gray-800 rounded-lg custom-scrollbar">
      <button
        onClick={handleCopy}
        className="absolute top-2 right-2 px-4 py-2 text-white rounded-lg hover:bg-blue-500 invisible group-hover:visible transition-all z-10"
        title="Copy code"
      >
        {copied ? (
          <Check className="w-4 h-4 text-green-400" />
        ) : (
          <Copy className="w-4 h-4 text-gray-300" />
        )}
      </button>
      <SyntaxHighlighter
        language={language}
        style={customStyle}
        showLineNumbers
        wrapLongLines
      >
        {codeContent}
      </SyntaxHighlighter>
    </div>
  );
};

interface ThinkingBlockProps {
  id: string;
  children: React.ReactNode;
  isExpanded: boolean;
  onToggle: (id: string) => void;
}

const ThinkingProcess: React.FC<ThinkingBlockProps> = ({ id, children, isExpanded, onToggle }) => {
  return (
    <div className="relative pt-1 pl-2">
      <button
        onClick={() => onToggle(id)}
        className="absolute -top-4 -left-6 p-1 bg-gray-100 hover:bg-gray-200 text-gray-500 rounded-full transition-colors z-20"
        title="View thinking process"
      >
        <Brain className="w-4 h-4" />
      </button>

      {isExpanded && (
        <div className="relative">
          <div className="absolute -left-4 top-0 flex flex-col gap-0.5">
            <div className="w-1.5 h-1.5 bg-gray-50 border border-gray-200 rounded-full ml-0.5" />
            <div className="w-2 h-2 bg-gray-50 border border-gray-200 rounded-full ml-1" />
            <div className="w-2.5 h-2.5 bg-gray-50 border border-gray-200 rounded-full ml-1.5" />
          </div>
          
          <div className="bg-gray-50 border border-gray-200 p-3 rounded-lg mt-4">
            <div className="text-sm text-gray-600">
              {children}
            </div>
          </div>
        </div>
      )}
    </div>
  );
};

const MathBlock = ({ math, isInParagraph }: { math: string; isInParagraph: boolean }) => {
  try {
    if (isInParagraph) {
      return (
        <span className="block my-2 overflow-x-auto text-center no-scrollbar">
          <InlineMath>{`\\displaystyle ${math}`}</InlineMath>
        </span>
      );
    }

    return (
      <div className="my-2 overflow-x-auto no-scrollbar">
        <BlockMath>{math}</BlockMath>
      </div>
    );
  } catch (error) {
    console.error('Math rendering error:', error);
    return <span className="text-red-500">Error rendering equation: {math}</span>;
  }
};

const ChatMessage = ({ 
  message, 
  onEdit, 
  isEditing, 
  editContent, 
  onEditChange, 
  onEditSave, 
  onEditCancel,
  setFullScreenImage
}: { 
  message: Message;
  onEdit?: (message: Message) => void;
  isEditing?: boolean;
  editContent?: string;
  onEditChange?: (content: string) => void;
  onEditSave?: () => void;
  onEditCancel?: () => void;
  setFullScreenImage: (image: string | null) => void;
}) => {
  const [expandedThinkBlocks, setExpandedThinkBlocks] = useState<Set<string>>(new Set());

  const toggleThinkBlock = (id: string) => {
    setExpandedThinkBlocks(prev => {
      const newSet = new Set(prev);
      if (newSet.has(id)) {
        newSet.delete(id);
      } else {
        newSet.add(id);
      }
      return newSet;
    });
  };

  const TextBlock = ({ children, isInParagraph = false }: { children: React.ReactNode; isInParagraph?: boolean }) => {
    if (Array.isArray(children)) {
      return <>{children.map((child, i) => <TextBlock key={i} isInParagraph={isInParagraph}>{child}</TextBlock>)}</>;
    }
  
    if (!children || typeof children !== 'string') {
      return <>{children}</>;
    }
  
    if (!children.includes('$')) {
      return <>{children}</>;
    }
  
    const segments: string[] = [];
    let currentText = '';
    let pos = 0;
    let inMath = false;
    let mathStart = 0;
  
    while (pos < children.length) {
      if (children[pos] === '$') {
        if (pos + 1 < children.length && children[pos + 1] === '$') {
          if (!inMath) {
            if (currentText) segments.push(currentText);
            currentText = '$$';
            mathStart = pos;
            inMath = true;
            pos += 2;
          } else if (pos > mathStart + 2) {
            segments.push(currentText + '$$');
            currentText = '';
            inMath = false;
            pos += 2;
          } else {
            currentText += '$$';
            pos += 2;
          }
        } else {
          if (!inMath) {
            if (currentText) segments.push(currentText);
            currentText = '$';
            mathStart = pos;
            inMath = true;
            pos++;
          } else if (pos > mathStart + 1) {
            segments.push(currentText + '$');
            currentText = '';
            inMath = false;
            pos++;
          } else {
            currentText += '$';
            pos++;
          }
        }
      } else {
        currentText += children[pos];
        pos++;
      }
    }
  
    if (currentText) {
      segments.push(currentText);
    }
  
    return (
      <>
        {segments.map((segment, i) => {
          try {
            if (segment.startsWith('$$') && segment.endsWith('$$')) {
              const math = segment.slice(2, -2);
              return <MathBlock key={i} math={math} isInParagraph={isInParagraph} />;
            } else if (segment.startsWith('$') && segment.endsWith('$')) {
              const math = segment.slice(1, -1);
              return <InlineMath key={i}>{math}</InlineMath>;
            }
            return <span key={i}>{segment}</span>;
          } catch (error) {
            console.error('LaTeX rendering error:', error);
            return <span key={i} className="text-red-500">{segment}</span>;
          }
        })}
      </>
    );
  };

  const components = {
    code: ({ children, inline, className }: { children: string; inline?: boolean; className?: string }) => (
      <CodeBlock inline={inline} className={className}>{children as string}</CodeBlock>
    ),
    think: ({ children }: { children: React.ReactNode }) => {
      const thinkId = useMemo(() => {
        const content = children?.toString() || '';
        return `think-${message.id}-${content.slice(0, 32)}`;
      }, [children, message.id]);

      return (
        <ThinkingProcess
          id={thinkId}
          isExpanded={expandedThinkBlocks.has(thinkId)}
          onToggle={toggleThinkBlock}
        >
          {children}
        </ThinkingProcess>
      );
    },
    p: ({ children }: { children: React.ReactNode }) => (
      <p className="my-1">
        <TextBlock isInParagraph={true}>{children}</TextBlock>
      </p>
    ),
    strong: ({ children }: { children: React.ReactNode }) => (
      <strong>
        <TextBlock>{children}</TextBlock>
      </strong>
    ),
    em: ({ children }: { children: React.ReactNode }) => (
      <em>
        <TextBlock>{children}</TextBlock>
      </em>
    ),
    li: ({ children }: { children: React.ReactNode }) => (
      <li>
        <TextBlock>{children}</TextBlock>
      </li>
    ),
    h1: ({ children }: { children: React.ReactNode }) => (
      <h1 className="text-2xl font-bold mb-4">
        <TextBlock>{children}</TextBlock>
      </h1>
    ),
    h2: ({ children }: { children: React.ReactNode }) => (
      <h2 className="text-xl font-bold mb-3">
        <TextBlock>{children}</TextBlock>
      </h2>
    ),
    h3: ({ children }: { children: React.ReactNode }) => (
      <h3 className="text-lg font-bold mb-2">
        <TextBlock>{children}</TextBlock>
      </h3>
    )
  };

  const processMessageContent = (content: string) => {
    if (content.includes('<think>') && !content.includes('</think>')) {
      content = content.replace(/<think>/g, '');
    }

    const thinkMatch = content.match(/<think>([\s\S]*?)<\/think>/);
    if (thinkMatch) {
      const thinking = thinkMatch[1].trim();
      const restOfContent = content.replace(/<think>[\s\S]*?<\/think>/, '').trim();
      return (
        <>
          <components.think>{thinking}</components.think>
          <ReactMarkdown 
            remarkPlugins={[remarkGfm]}
            components={components as Partial<Components>}
          >
            {restOfContent}
          </ReactMarkdown>
        </>
      );
    }
    return (
      <ReactMarkdown 
        remarkPlugins={[remarkGfm]}
        components={components as Partial<Components>}
      >
        {content}
      </ReactMarkdown>
    );
  };

  const downloadImage = async (imageUrl: string) => {
    try {
      const response = await fetch(imageUrl);
      const blob = await response.blob();
      const url = window.URL.createObjectURL(blob);
      const link = document.createElement('a');
      link.href = url;
      link.download = `image-${Date.now()}.png`;
      document.body.appendChild(link);
      link.click();
      document.body.removeChild(link);
      window.URL.revokeObjectURL(url);
    } catch (error) {
      console.error('Error downloading image:', error);
    }
  };

  return (
    <div className={`flex ${message.role === 'user' ? 'justify-end' : 'justify-start'}`}>
      <div
        className={`relative max-w-[80%] md:max-w-[70%] lg:max-w-[60%] rounded-2xl px-4 py-2 \
          ${message.role === 'user'
            ? 'bg-blue-500 text-white rounded-br-none dark:bg-blue-600'
            : 'bg-white text-gray-800 rounded-bl-none shadow-sm dark:bg-gray-800 dark:text-gray-100 dark:shadow-none'}
        `}
      >
        {isEditing ? (
          <div className="flex flex-col gap-2">
            <textarea
              value={editContent}
              onChange={(e) => onEditChange?.(e.target.value)}
              className="w-full p-2 rounded border-none bg-blue-500 text-white placeholder-white/75 focus:outline-none custom-scrollbar"
              rows={5}
              cols={50}
              placeholder="Edit your message..."
            />
            <div className="flex justify-end gap-2">
              <button
                onClick={onEditCancel}
                className="p-1.5 rounded-full bg-blue-600 hover:bg-blue-700 transition-colors"
                title="Cancel"
              >
                <X className="w-4 h-4" />
              </button>
              <button
                onClick={onEditSave}
                className="p-1.5 rounded-full bg-blue-600 hover:bg-blue-700 transition-colors"
                title="Save changes"
              >
                <Check className="w-4 h-4" />
              </button>
            </div>
          </div>
        ) : (
          <>
            {message.image && (
              <div className="relative group">
                <img 
                  src={message.image.startsWith('https://image.pollinations.ai') ? '/loading.gif' : message.image} 
                  alt="Attached" 
                  className="max-w-full rounded-lg mb-2 max-h-[300px] object-contain cursor-pointer"
                  onClick={() => setFullScreenImage(message.image || null)}
                  onLoad={(e) => {
                    if (message.image?.startsWith('https://image.pollinations.ai')) {
                      e.currentTarget.src = message.image!;
                    }
                  }}
                />
                <button
                  onClick={() => downloadImage(message.image!)}
                  className="absolute top-2 right-2 p-2 bg-black/50 text-white rounded-full opacity-0 group-hover:opacity-100 transition-opacity hover:bg-black/70"
                  title="Download image"
                >
                  <Download className="w-4 h-4" />
                </button>
              </div>
            )}
            <div className={`prose ${message.role === 'user' ? 'prose-invert' : ''} max-w-none prose-sm prose-p:my-1 prose-pre:my-1 prose-ul:my-1 prose-ol:my-1 prose-li:my-0 dark:prose-invert`}>
              {processMessageContent(message.content)}
            </div>
            <div className="text-xs opacity-70 mt-1 text-right flex items-center justify-end gap-2">
              {message.role === 'user' && (
                <button
                  onClick={() => onEdit?.(message)}
                  className="p-1 hover:bg-white/20 rounded transition-colors"
                  title="Edit message"
                >
                  <Pencil className="w-3 h-3" />
                </button>
              )}
              {message.timestamp.toLocaleTimeString([], { 
                hour: '2-digit', 
                minute: '2-digit' 
              })}
            </div>
          </>
        )}
      </div>
    </div>
  );
};


function App() {
  const [messages, setMessages] = useState<Message[]>(() => {
    const savedMessages = localStorage.getItem(STORAGE_KEY);
    if (savedMessages) {
      const parsed = JSON.parse(savedMessages);
      return parsed.map((msg: any) => ({
        ...msg,
        timestamp: new Date(msg.timestamp)
      }));
    }
    return [SYSTEM_PROMPT];
  });
  
  const [input, setInput] = useState('');
  const [isLoading, setIsLoading] = useState(false);
  const [showClearConfirm, setShowClearConfirm] = useState(false);
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const [selectedImage, setSelectedImage] = useState<string | null>(null);
  const fileInputRef = useRef<HTMLInputElement>(null);
  const [fullScreenImage, setFullScreenImage] = useState<string | null>(null);
  const textareaRef = useRef<HTMLTextAreaElement>(null);
  const [reasonerEnabled, setReasonerEnabled] = useState(false);
  const [editingMessage, setEditingMessage] = useState<EditingMessage | null>(null);
  const [searchEnabled, setSearchEnabled] = useState(false);
  const [imageGenEnabled, setImageGenEnabled] = useState(false);
  const [generatingImage, setGeneratingImage] = useState(false);
  const [darkMode, setDarkMode] = useState(() => {
    if (typeof window !== 'undefined') {
      return localStorage.getItem('theme') === 'dark' ||
        (window.matchMedia && window.matchMedia('(prefers-color-scheme: dark)').matches);
    }
    return false;
  });

  useEffect(() => {
    if (darkMode) {
      document.documentElement.classList.add('dark');
      localStorage.setItem('theme', 'dark');
    } else {
      document.documentElement.classList.remove('dark');
      localStorage.setItem('theme', 'light');
    }
  }, [darkMode]);

  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'auto' });
  };

  useEffect(() => {
    scrollToBottom();
  }, [messages]);

  // Save messages to localStorage whenever they change
  useEffect(() => {
    localStorage.setItem(STORAGE_KEY, JSON.stringify(messages));
  }, [messages]);

  const handleClearHistory = () => {
    setMessages([SYSTEM_PROMPT as Message]);
    setShowClearConfirm(false);
  };

  const toggleMode = (mode: 'reasoner' | 'search' | 'image' | 'generate') => {
    switch (mode) {
      case 'reasoner':
        if (reasonerEnabled) {
          setReasonerEnabled(false);
        } else {
          setReasonerEnabled(true);
          setSearchEnabled(false);
          setImageGenEnabled(false);
          setSelectedImage(null);
          if (fileInputRef.current) {
            fileInputRef.current.value = '';
          }
        }
        break;
      case 'search':
        if (searchEnabled) {
          setSearchEnabled(false);
        } else {
          setSearchEnabled(true);
          setReasonerEnabled(false);
          setImageGenEnabled(false);
          setSelectedImage(null);
          if (fileInputRef.current) {
            fileInputRef.current.value = '';
          }
        }
        break;
      case 'image':
        if (selectedImage) {
          setSelectedImage(null);
          if (fileInputRef.current) {
            fileInputRef.current.value = '';
          }
        } else {
          fileInputRef.current?.click();
          setReasonerEnabled(false);
          setSearchEnabled(false);
          setImageGenEnabled(false);
        }
        break;
      case 'generate':
        if (imageGenEnabled) {
          setImageGenEnabled(false);
        } else {
          setImageGenEnabled(true);
          setReasonerEnabled(false);
          setSearchEnabled(false);
          setSelectedImage(null);
          if (fileInputRef.current) {
            fileInputRef.current.value = '';
          }
        }
        break;
    }
  };

  const generateImage = async (prompt: string) => {
    if (!prompt.trim() || generatingImage) return;

    setGeneratingImage(true);
    try {
      const encodedPrompt = encodeURIComponent(prompt.trim());
      const imageUrl = `https://image.pollinations.ai/prompt/${encodedPrompt}?width=1024&height=1024&nologo=true&private=true&enhance=true&safe=true`;
      
      const userMessage: Message = {
        id: Date.now().toString(),
        role: 'user',
        content: `Generate an image: ${prompt}`,
        timestamp: new Date()
      };

      setMessages(prev => [...prev, userMessage]);

      const assistantMessage: Message = {
        id: (Date.now() + 1).toString(),
        role: 'assistant',
        content: 'Here\'s your generated image!',
        timestamp: new Date(),
        image: imageUrl
      };

      setMessages(prev => [...prev, assistantMessage]);
      setInput('');
    } catch (error) {
      console.error('Error generating image:', error);
      setMessages(prev => [...prev, {
        id: Date.now().toString(),
        role: 'assistant',
        content: 'Sorry, I encountered an error while generating the image. Please try again.',
        timestamp: new Date()
      }]);
    } finally {
      setGeneratingImage(false);
    }
  };

  const handleImageSelect = (e: React.ChangeEvent<HTMLInputElement>) => {
    const file = e.target.files?.[0];
    if (file) {
      if (file.size > 5 * 1024 * 1024) { // 5MB limit
        alert('Image size should be less than 5MB');
        return;
      }

      const reader = new FileReader();
      reader.onloadend = () => {
        setSelectedImage(reader.result as string);
        setReasonerEnabled(false);
        setSearchEnabled(false);
      };
      reader.readAsDataURL(file);
    }
  };

  const removeSelectedImage = () => {
    setSelectedImage(null);
    if (fileInputRef.current) {
      fileInputRef.current.value = '';
    }
  };

  const makeApiRequest = async (messages: Message[], retryCount = 0, useFallback = false): Promise<any> => {
    const randomApiKey = GROQ_API_KEYS[Math.floor(Math.random() * GROQ_API_KEYS.length)];
    
    try {
      const systemPrompt = messages.find(msg => msg.role === 'system');
      const lastMessages = messages.filter(msg => msg.role !== 'system').slice(-5);
      const apiMessages = systemPrompt ? [systemPrompt, ...lastMessages] : lastMessages;

      const hasImageInHistory = apiMessages.some(msg => !!msg.image);
      const filteredMessages = apiMessages.map(msg => {
        // Skip messages with images if using reasoning model
        if (reasonerEnabled && msg.image) {
          return null;
        }
        if (msg.role === 'assistant' && msg.image) {
          return null;
        }

        const lastImageMessage = hasImageInHistory ? 
          [...apiMessages].reverse().find((m: Message) => !!m.image) : null;

        if (msg.image && msg === lastImageMessage && !reasonerEnabled) {
          return {
            role: msg.role,
            content: [
              { type: "text", text: msg.content },
              { type: "image_url", image_url: { url: msg.image } }
            ]
          };
        } else {
          return {
            role: msg.role,
            content: msg.content
          };
        }
      }).filter(Boolean); // Remove any null messages (filtered image messages)

      let model;
      let apiUrl;

      if (reasonerEnabled) {
        // Use GROQ API for reasoning
        model = GROQ_MODEL_REASONER;
        apiUrl = GROQ_API_URL;
      } else {
        // Use Pollinations API for everything else
        model = searchEnabled ? POLLINATIONS_MODEL_SEARCH :
          (hasImageInHistory ? POLLINATIONS_MODEL_VISION :
            (useFallback ? POLLINATIONS_MODEL_FALLBACK : POLLINATIONS_MODEL_TEXT));
        apiUrl = POLLINATIONS_API_URL;
      }

      // Create AbortController for timeout
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT);

      try {
        const headers: Record<string, string> = {
          'Content-Type': 'application/json'
        };

        if (apiUrl === GROQ_API_URL) {
          headers['Authorization'] = `Bearer ${randomApiKey}`;
        }

        const response = await fetch(apiUrl, {
          method: 'POST',
          headers,
          body: JSON.stringify({
            model,
            messages: filteredMessages,
            temperature: 0.7,
            max_tokens: (model === GROQ_MODEL_REASONER) ? 8192 : 4096,
            stream: true
          }),
          signal: controller.signal
        });

        clearTimeout(timeoutId); // Clear timeout if request completes

        if (!response.ok) {
          const errorData = await response.json().catch(() => ({}));
          throw new Error(`HTTP error! status: ${response.status}, message: ${errorData.error?.message || 'Unknown error'}`);
        }

        return response;

      } catch (error) {
        clearTimeout(timeoutId); // Clear timeout on error
        if (error instanceof Error && error.name === 'AbortError') {
          throw new Error('Request timed out. Please try again.');
        }
        throw error;
      }

    } catch (error) {
      // If error is timeout and not using fallback, try fallback immediately
      if (error instanceof Error && error.message === 'Request timed out. Please try again.' && !useFallback) {
        console.log('Request timed out, switching to fallback model...');
        return makeApiRequest(messages, 0, true);
      }

      if (!useFallback) {
        console.log('Switching to fallback model...');
        return makeApiRequest(messages, 0, true);
      }
      
      if (retryCount < MAX_RETRIES) {
        console.log(`Retry attempt ${retryCount + 1} of ${MAX_RETRIES}`);
        await new Promise(resolve => setTimeout(resolve, RETRY_DELAY * (retryCount + 1)));
        return makeApiRequest(messages, retryCount + 1, true);
      }
      
      throw error;
    }
  };

  const fetchYouTubeTranscript = async (url: string): Promise<string> => {
    try {
      const response = await fetch(TACTIQ_API_URL, {
        method: 'POST',
        headers: {
          'accept': '*/*',
          'content-type': 'application/json',
          'sec-ch-ua': '"Chromium";v="134", "Not:A-Brand";v="24", "Brave";v="134"',
          'sec-ch-ua-mobile': '?0',
          'sec-ch-ua-platform': '"Windows"',
          'Referer': 'https://tactiq.io/',
        },
        body: JSON.stringify({ videoUrl: url })
      });

      if (!response.ok) {
        throw new Error('Failed to fetch transcript');
      }

      const data = await response.json();
      let text = `Title: ${data.title}\n\nTranscript:\n`;
      data.captions.forEach((caption: any) => {
        text += caption.text + ' ';
      });

      return text.trim();
    } catch (error) {
      console.error('Error fetching YouTube transcript:', error);
      throw error;
    }
  };

  const extractYouTubeUrl = (text: string): string | null => {
    const regex = /(?:https?:\/\/)?(?:www\.)?(?:youtube\.com\/watch\?v=|youtu\.be\/)([^\s&]+)/;
    const match = text.match(regex);
    return match ? match[0] : null;
  };

  const handleSubmit = async (e: React.FormEvent) => {
    e.preventDefault();
    if ((!input.trim() && !selectedImage) || isLoading) return;

    // If image generation is enabled, use the generateImage function instead
    if (imageGenEnabled || input.trim().startsWith('Generate an image')) {
      const prompt = input.trim().replace('Generate an image: ', '');
      const reprompt = prompt.replace('Generate an image', '');
      await generateImage(reprompt.trim());
      return;
    }

    // Create visible user message
    const userMessage: Message = {
      id: Date.now().toString(),
      role: 'user',
      content: input.trim() || "What's in this image?",
      timestamp: new Date(),
      image: selectedImage || undefined
    };

    setMessages(prev => [...prev, userMessage]);
    setInput('');
    setSelectedImage(null);
    if (fileInputRef.current) {
      fileInputRef.current.value = '';
    }
    setIsLoading(true);

    if (textareaRef.current) {
      textareaRef.current.style.height = 'auto';
      textareaRef.current.focus();
    }

    try {
      // Check for YouTube URL and get transcript if exists
      const youtubeUrl = extractYouTubeUrl(userMessage.content);
      let transcriptText = '';
      if (youtubeUrl) {
        try {
          transcriptText = await fetchYouTubeTranscript(youtubeUrl);
        } catch (error) {
          console.error('Failed to fetch YouTube transcript:', error);
        }
      }

      const apiMessages: Message[] = selectedImage ? 
        [{ ...SYSTEM_PROMPT, role: 'system' as const }] : 
        [...messages];

      apiMessages.push({
        ...userMessage,
        content: transcriptText ? 
          `${userMessage.content}\n\nContext from video:\n${transcriptText}` : 
          userMessage.content
      });

      const response = await makeApiRequest(apiMessages);
      const reader = response.body?.getReader();
      if (!reader) throw new Error('Failed to get response reader');

      // Create a new message for streaming response
      const assistantMessage: Message = {
        id: Date.now().toString(),
        role: 'assistant',
        content: '',
        timestamp: new Date(),
      };
      setMessages(prev => [...prev, assistantMessage]);

      // Read the stream
      const decoder = new TextDecoder();
      let buffer = '';
      
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        
        const lines = buffer.split('\n');
        buffer = lines.pop() || ''; // Keep the incomplete line in buffer

        for (const line of lines) {
          if (line.trim() === '') continue;
          if (line === 'data: [DONE]') continue;
          
          try {
            const data = JSON.parse(line.replace(/^data: /, ''));
            const content = data.choices[0]?.delta?.content || '';
            
            setMessages(prev => prev.map(msg =>
              msg.id === assistantMessage.id
                ? { ...msg, content: msg.content + content }
                : msg
            ));
          } catch (error) {
            console.error('Error parsing stream:', error);
          }
        }
      }

    } catch (error) {
      console.error('Error:', error);
      setMessages(prev => [...prev, {
        id: Date.now().toString(),
        role: 'assistant',
        content: 'Sorry, I encountered an error. Please try again. If the problem persists, try refreshing the page or clearing the chat history.',
        timestamp: new Date(),
      }]);
    } finally {
      setIsLoading(false);
      requestAnimationFrame(() => {
        if (textareaRef.current) {
          textareaRef.current.focus();
        }
      });
    }
  };

  const handleEditMessage = (message: Message) => {
    setEditingMessage({
      id: message.id,
      content: message.content,
      image: message.image
    });
  };

  const handleEditCancel = () => {
    setEditingMessage(null);
  };

  const handleEditSave = async () => {
    if (!editingMessage) return;

    const messageIndex = messages.findIndex(m => m.id === editingMessage.id);
    if (messageIndex === -1) return;

    // Check if content has actually changed
    const originalMessage = messages[messageIndex];
    if (originalMessage.content === editingMessage.content && 
        originalMessage.image === editingMessage.image) {
      setEditingMessage(null);
      return;
    }

    const updatedMessages = [...messages];
    updatedMessages[messageIndex] = {
      ...updatedMessages[messageIndex],
      content: editingMessage.content,
      image: editingMessage.image
    };

    updatedMessages.splice(messageIndex + 1);
    setMessages(updatedMessages);
    setEditingMessage(null);

    // If the edited message was an image generation prompt, regenerate the image
    if (imageGenEnabled || editingMessage.content.trim().startsWith('Generate an image')) {
      const prompt = editingMessage.content.trim().replace('Generate an image: ', '');
      const reprompt = prompt.replace('Generate an image', '');
      const encodedPrompt = encodeURIComponent(reprompt.trim());
      const imageUrl = `https://image.pollinations.ai/prompt/${encodedPrompt}?width=1024&height=1024&nologo=true&private=true&enhance=true&safe=true`;
      const assistantMessage: Message = {
        id: (Date.now() + 1).toString(),
        role: 'assistant',
        content: 'Here\'s your generated image!',
        timestamp: new Date(),
        image: imageUrl
      };

      setMessages(prev => [...prev, assistantMessage]);
      return;
    }

    setIsLoading(true);
    try {
      const response = await makeApiRequest(updatedMessages);
      const reader = response.body?.getReader();
      if (!reader) throw new Error('Failed to get response reader');

      const assistantMessage: Message = {
        id: Date.now().toString(),
        role: 'assistant',
        content: '',
        timestamp: new Date(),
      };
      setMessages(prev => [...prev, assistantMessage]);

      const decoder = new TextDecoder();
      let buffer = '';
      
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';

        for (const line of lines) {
          if (line.trim() === '' || line === 'data: [DONE]') continue;
          
          try {
            const data = JSON.parse(line.replace(/^data: /, ''));
            const content = data.choices[0]?.delta?.content || '';
            
            setMessages(prev => prev.map(msg =>
              msg.id === assistantMessage.id
                ? { ...msg, content: msg.content + content }
                : msg
            ));
          } catch (error) {
            console.error('Error parsing stream:', error);
          }
        }
      }
    } catch (error) {
      console.error('Error:', error);
      setMessages(prev => [...prev, {
        id: Date.now().toString(),
        role: 'assistant',
        content: 'Sorry, I encountered an error. Please try again. If the problem persists, try refreshing the page or clearing the chat history.',
        timestamp: new Date(),
      }]);
    } finally {
      setIsLoading(false);
    }
  };

  // Add a warning if messages fail to load from localStorage
  useEffect(() => {
    const savedMessages = localStorage.getItem(STORAGE_KEY);
    if (savedMessages) {
      try {
        JSON.parse(savedMessages);
      } catch (error) {
        console.error('Error loading chat history:', error);
        localStorage.removeItem(STORAGE_KEY);
        setMessages([SYSTEM_PROMPT as Message]);
      }
    }
  }, []);

  // Add clipboard paste handler
  const handlePaste = async (e: ClipboardEvent) => {
    const items = e.clipboardData?.items;
    if (!items) return;

    for (const item of Array.from(items)) {
      if (item.type.startsWith('image/')) {
        e.preventDefault();
        const file = item.getAsFile();
        if (!file) continue;

        if (file.size > 5 * 1024 * 1024) { // 5MB limit
          alert('Image size should be less than 5MB');
          return;
        }

        const reader = new FileReader();
        reader.onloadend = () => {
          setSelectedImage(reader.result as string);
          setReasonerEnabled(false);
          setSearchEnabled(false);
        };
        reader.readAsDataURL(file);
        break;
      }
    }
  };

  useEffect(() => {
    document.addEventListener('paste', handlePaste);
    return () => {
      document.removeEventListener('paste', handlePaste);
    };
  }, []);

  const adjustTextareaHeight = () => {
    const textarea = textareaRef.current;
    if (textarea) {
      textarea.style.height = 'auto';
      textarea.style.height = `${Math.min(textarea.scrollHeight, 200)}px`; // Max height of 200px
    }
  };

  const handleInputChange = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
    setInput(e.target.value);
    adjustTextareaHeight();
  };

  useEffect(() => {
    if (textareaRef.current) {
      textareaRef.current.focus();
    }
  }, []);

  return (
    <div className="fixed inset-0 bg-gray-100 dark:bg-gray-900 z-50 flex flex-col">
      {/* Header */}
      <div className="bg-white dark:bg-gray-900 shadow-sm dark:shadow-none flex-shrink-0 border-b border-gray-200 dark:border-gray-800">
        <div className="w-full px-2 sm:px-4 flex items-center justify-between py-4">
          <div className="flex items-center gap-3">
            <div className="flex items-center gap-2">
              <div className="w-10 h-10 bg-blue-500 rounded-full flex items-center justify-center">
                <Bot className="w-6 h-6 text-white" />
              </div>
              <div>
                <h1 className="font-semibold text-gray-900 dark:text-gray-100">Study Assistant</h1>
                <p className="text-sm text-gray-500 dark:text-green-400">
                  {isLoading ? 'Typing...' : 'Online'}
                </p>
              </div>
            </div>
          </div>
          <div className="flex items-center gap-2">
            <button
              onClick={() => setDarkMode((d) => !d)}
              className="p-2 rounded-lg hover:bg-gray-100 dark:hover:bg-gray-800 transition-colors"
              title={darkMode ? 'Switch to light mode' : 'Switch to dark mode'}
            >
              {darkMode ? <Sun className="w-5 h-5 text-yellow-400" /> : <Moon className="w-5 h-5 text-gray-700" />}
            </button>
            <button
              onClick={() => setShowClearConfirm(true)}
              className="p-2 hover:bg-gray-100 dark:hover:bg-gray-800 rounded-lg transition-colors text-red-600"
              title="Clear chat history"
            >
              <Trash2 className="w-5 h-5" />
            </button>
          </div>
        </div>
      </div>

      {/* Messages */}
      <div className="flex-1 overflow-y-auto min-h-0 no-scrollbar">
        <div className="w-full px-2 sm:px-8">
          {messages.length <= 1 ? (
            <div className="flex justify-center items-center min-h-[80vh]">
              <div className="bg-white dark:bg-gray-800 rounded-xl shadow-lg p-8 max-w-xl w-full text-center border border-gray-200 dark:border-gray-700">
                <div className="flex flex-col items-center mb-4">
                  <span className="text-4xl mb-2 animate-bounce">👋</span>
                  <h2 className="text-2xl font-bold mb-1 text-gray-900 dark:text-gray-100">Welcome to Study Assistant!</h2>
                </div>
                <p className="mb-4 text-gray-600 dark:text-gray-300 text-lg">
                  I'm your friendly AI study buddy—here to help you learn, solve problems, and master new topics. Just ask me anything!
                </p>
                <ul className="text-left text-gray-700 dark:text-gray-200 list-disc list-inside space-y-2 mb-4">
                  <li><b>Breaks down</b> tough topics into simple, bite-sized explanations.</li>
                  <li><b>Guides you</b> step-by-step through problems and concepts.</li>
                  <li><b>Encourages</b> critical thinking and independent learning.</li>
                  <li><b>Supports you</b> with tips, resources, and clear answers.</li>
                </ul>
                <div className="mt-4 text-gray-500 dark:text-gray-400 text-base">
                  🚀 <b>Get started:</b> Type your question below and let's learn together!
                </div>
                <div className="mt-8 text-xs text-gray-400 dark:text-gray-500">
                  Developed by <a href="https://www.instagram.com/sanch1t_" target="_blank" rel="noopener noreferrer" className="underline hover:text-blue-500">Sanchit</a>
                </div>
              </div>
            </div>
          ) : (
            <div className="py-4 space-y-4">
              {messages.slice(1).map(message => (
                <ChatMessage
                  key={message.id}
                  message={message}
                  onEdit={handleEditMessage}
                  isEditing={editingMessage?.id === message.id}
                  editContent={editingMessage?.content}
                  onEditChange={(content) => setEditingMessage(prev => prev ? { ...prev, content } : null)}
                  onEditSave={handleEditSave}
                  onEditCancel={handleEditCancel}
                  setFullScreenImage={setFullScreenImage}
                />
              ))}
              {isLoading && (
                <div className="flex justify-start">
                  <div className="bg-white text-gray-800 rounded-2xl rounded-bl-none px-4 py-2 shadow-sm dark:bg-gray-800 dark:text-gray-100 dark:shadow-none">
                    <div className="flex space-x-2">
                      <div className="w-2 h-2 rounded-full animate-bounce bg-gray-400 dark:bg-gray-300" />
                      <div className="w-2 h-2 rounded-full animate-bounce bg-gray-400 dark:bg-gray-300" style={{ animationDelay: '0.2s' }} />
                      <div className="w-2 h-2 rounded-full animate-bounce bg-gray-400 dark:bg-gray-300" style={{ animationDelay: '0.4s' }} />
                    </div>
                  </div>
                </div>
              )}
              <div ref={messagesEndRef} />
            </div>
          )}
        </div>
      </div>

      {/* Input section */}
      <div className="bg-white dark:bg-gray-900 shadow-lg dark:shadow-none mt-auto flex-shrink-0 border-t border-gray-200 dark:border-gray-800">
        <div className="w-full px-2 sm:px-4">
          <div className="py-4">
            {selectedImage && (
              <div className="mb-2 relative inline-block">
                <img 
                  src={selectedImage} 
                  alt="Selected" 
                  className="h-20 rounded-lg object-contain"
                />
                <button
                  onClick={removeSelectedImage}
                  className="absolute -top-2 -right-2 p-1 bg-red-500 text-white rounded-full hover:bg-red-600 transition-colors"
                >
                  <X className="w-4 h-4" />
                </button>
              </div>
            )}
            <form onSubmit={handleSubmit} className="flex gap-4 items-start max-w-full">
              <div className="flex gap-1">
                <button
                  type="button"
                  onClick={() => toggleMode('reasoner')}
                  className={`p-2 rounded-full transition-colors flex-shrink-0 ${
                    reasonerEnabled 
                      ? 'text-white bg-blue-500 hover:bg-blue-600' 
                      : 'text-gray-500 hover:bg-gray-100 hover:text-gray-700'
                  }`}
                  title={`${reasonerEnabled ? 'Disable' : 'Enable'} reasoner mode`}
                >
                  <Brain className="w-5 h-5" />
                </button>
                <button
                  type="button"
                  onClick={() => toggleMode('search')}
                  className={`p-2 rounded-full transition-colors flex-shrink-0 ${
                    searchEnabled 
                      ? 'text-white bg-blue-500 hover:bg-blue-600' 
                      : 'text-gray-500 hover:bg-gray-100 hover:text-gray-700'
                  }`}
                  title={`${searchEnabled ? 'Disable' : 'Enable'} search mode`}
                >
                  <Globe className="w-5 h-5" />
                </button>
                <button
                  type="button"
                  onClick={() => toggleMode('generate')}
                  className={`p-2 rounded-full transition-colors flex-shrink-0 ${
                    imageGenEnabled 
                      ? 'text-white bg-blue-500 hover:bg-blue-600' 
                      : 'text-gray-500 hover:bg-gray-100 hover:text-gray-700'
                  }`}
                  title={`${imageGenEnabled ? 'Disable' : 'Enable'} image generation mode`}
                  disabled={generatingImage}
                >
                  <Wand2 className="w-5 h-5" />
                </button>
                <button
                  type="button"
                  onClick={() => toggleMode('image')}
                  disabled={isLoading}
                  className={`group relative p-2 rounded-full transition-colors flex-shrink-0 ${
                    selectedImage 
                      ? 'text-white bg-blue-500 hover:bg-blue-600' 
                      : 'text-gray-500 hover:bg-gray-100 hover:text-gray-700'
                  }`}
                  title="Attach image or paste from clipboard"
                >
                  <ImageUp className="w-5 h-5" />
                </button>
              </div>
              <textarea
                ref={textareaRef}
                value={input}
                onChange={handleInputChange}
                onKeyDown={(e) => {
                  if (e.key === 'Enter' && !e.shiftKey) {
                    e.preventDefault();
                    handleSubmit(e);
                  }
                }}
                placeholder={selectedImage ? "Ask about this image..." : "Ask me anything..."}
                className="flex-1 px-4 py-2 rounded-lg border bg-transparent focus:outline-none resize-none min-h-[40px] max-h-[100px] overflow-y-auto custom-scrollbar text-gray-900 dark:text-gray-100 placeholder-gray-400 dark:placeholder-gray-500 border-gray-300 dark:border-gray-700 bg-white dark:bg-gray-800"
                disabled={isLoading}
                rows={1}
              />
              <input
                type="file"
                accept="image/*"
                onChange={handleImageSelect}
                ref={fileInputRef}
                className="hidden"
              />
              <button
                type="submit"
                disabled={(!input.trim() && !selectedImage) || isLoading || generatingImage}
                className="w-10 h-10 bg-blue-500 text-white rounded-full hover:bg-blue-600 disabled:opacity-50 disabled:cursor-not-allowed transition-colors flex items-center justify-center flex-shrink-0"
              >
                <ArrowUp className="w-5 h-5" />
              </button>
            </form>
          </div>
        </div>
      </div>

      {showClearConfirm && (
        <div className="fixed inset-0 bg-black/50 flex items-center justify-center z-50">
          <div className="bg-white dark:bg-gray-800 rounded-lg p-6 max-w-sm w-full mx-4">
            <h3 className="text-lg font-semibold mb-2">Clear Chat History</h3>
            <p className="text-gray-600 mb-4">
              Are you sure you want to clear all chat history? This action cannot be undone.
            </p>
            <div className="flex justify-end gap-2">
              <button
                onClick={() => setShowClearConfirm(false)}
                className="px-4 py-2 text-gray-600 hover:bg-gray-100 rounded-lg transition-colors"
              >
                Cancel
              </button>
              <button
                onClick={handleClearHistory}
                className="px-4 py-2 bg-red-500 text-white rounded-lg hover:bg-red-600 transition-colors"
              >
                Clear History
              </button>
            </div>
          </div>
        </div>
      )}

      {fullScreenImage && (
        <div 
          className="fixed inset-0 bg-black/90 z-50 flex items-center justify-center cursor-pointer"
          onClick={() => setFullScreenImage(null)}
        >
          <button
            onClick={() => setFullScreenImage(null)}
            className="absolute top-4 right-4 p-2 bg-black/50 text-white rounded-full hover:bg-black/70 transition-colors"
          >
            <X className="w-6 h-6" />
          </button>
          <img 
            src={fullScreenImage} 
            alt="Full screen"
            className="max-w-[90%] max-h-[90vh] object-contain"
            onClick={(e) => e.stopPropagation()}
          />
        </div>
      )}
    </div>
  );
};


export default App