Atomizer/atomizer-dashboard/frontend/src/components/ClaudeChat.tsx

import React, { useState, useRef, useEffect } from 'react';
import {
  Send,
  Bot,
  User,
  Sparkles,
  Loader2,
  X,
  Maximize2,
  Minimize2,
  AlertCircle,
  Wrench,
  ChevronDown,
  ChevronUp,
  Trash2
} from 'lucide-react';
import ReactMarkdown from 'react-markdown';
import remarkGfm from 'remark-gfm';
import { useStudy } from '../context/StudyContext';

interface Message {
  id: string;
  role: 'user' | 'assistant';
  content: string;
  timestamp: Date;
  toolCalls?: Array<{
    tool: string;
    input: Record<string, any>;
    result_preview: string;
  }>;
}

interface ClaudeChatProps {
  isExpanded?: boolean;
  onToggleExpand?: () => void;
  onClose?: () => void;
}

export const ClaudeChat: React.FC<ClaudeChatProps> = ({
  isExpanded = false,
  onToggleExpand,
  onClose
}) => {
  const { selectedStudy } = useStudy();
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState('');
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [apiAvailable, setApiAvailable] = useState<boolean | null>(null);
  const [suggestions, setSuggestions] = useState<string[]>([]);
  const [expandedTools, setExpandedTools] = useState<Set<string>>(new Set());
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const inputRef = useRef<HTMLInputElement>(null);

  // Check API status on mount
  useEffect(() => {
    checkApiStatus();
  }, []);

  // Load suggestions when study changes
  useEffect(() => {
    loadSuggestions();
  }, [selectedStudy]);

  // Scroll to bottom when messages change
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages]);

  const checkApiStatus = async () => {
    try {
      const response = await fetch('/api/claude/status');
      const data = await response.json();
      setApiAvailable(data.available);
      if (!data.available) {
        setError(data.message);
      }
    } catch (err) {
      setApiAvailable(false);
      setError('Could not connect to Claude API');
    }
  };

  const loadSuggestions = async () => {
    try {
      const url = selectedStudy
        ? `/api/claude/suggestions?study_id=${selectedStudy.id}`
        : '/api/claude/suggestions';
      const response = await fetch(url);
      const data = await response.json();
      setSuggestions(data.suggestions || []);
    } catch (err) {
      setSuggestions([]);
    }
  };

  const sendMessage = async (messageText?: string) => {
    const text = messageText || input.trim();
    if (!text || isLoading) return;

    setError(null);
    const userMessage: Message = {
      id: `user-${Date.now()}`,
      role: 'user',
      content: text,
      timestamp: new Date()
    };

    setMessages(prev => [...prev, userMessage]);
    setInput('');
    setIsLoading(true);

    try {
      const response = await fetch('/api/claude/chat', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          message: text,
          study_id: selectedStudy?.id,
          conversation_history: messages.map(m => ({
            role: m.role,
            content: m.content
          }))
        })
      });

      if (!response.ok) {
        const errorData = await response.json();
        throw new Error(errorData.detail || 'Failed to get response');
      }

      const data = await response.json();

      const assistantMessage: Message = {
        id: `assistant-${Date.now()}`,
        role: 'assistant',
        content: data.response,
        timestamp: new Date(),
        toolCalls: data.tool_calls
      };

      setMessages(prev => [...prev, assistantMessage]);
    } catch (err: any) {
      setError(err.message || 'Failed to send message');
    } finally {
      setIsLoading(false);
      inputRef.current?.focus();
    }
  };

  const clearConversation = () => {
    setMessages([]);
    setError(null);
  };

  const toggleToolExpand = (toolId: string) => {
    setExpandedTools(prev => {
      const newSet = new Set(prev);
      if (newSet.has(toolId)) {
        newSet.delete(toolId);
      } else {
        newSet.add(toolId);
      }
      return newSet;
    });
  };

  // Render tool call indicator
  const renderToolCalls = (toolCalls: Message['toolCalls'], messageId: string) => {
    if (!toolCalls || toolCalls.length === 0) return null;

    return (
      <div className="mt-2 space-y-1">
        {toolCalls.map((tool, index) => {
          const toolId = `${messageId}-tool-${index}`;
          const isExpanded = expandedTools.has(toolId);

          return (
            <div
              key={index}
              className="bg-dark-700/50 rounded-lg border border-dark-600 overflow-hidden"
            >
              <button
                onClick={() => toggleToolExpand(toolId)}
                className="w-full px-3 py-2 flex items-center justify-between text-xs hover:bg-dark-600/50 transition-colors"
              >
                <div className="flex items-center gap-2">
                  <Wrench className="w-3 h-3 text-primary-400" />
                  <span className="text-dark-300">Used tool: </span>
                  <span className="text-primary-400 font-mono">{tool.tool}</span>
                </div>
                {isExpanded ? (
                  <ChevronUp className="w-3 h-3 text-dark-400" />
                ) : (
                  <ChevronDown className="w-3 h-3 text-dark-400" />
                )}
              </button>
              {isExpanded && (
                <div className="px-3 py-2 border-t border-dark-600 text-xs">
                  <div className="text-dark-400 mb-1">Input:</div>
                  <pre className="text-dark-300 bg-dark-800 p-2 rounded overflow-x-auto">
                    {JSON.stringify(tool.input, null, 2)}
                  </pre>
                  <div className="text-dark-400 mt-2 mb-1">Result preview:</div>
                  <pre className="text-dark-300 bg-dark-800 p-2 rounded overflow-x-auto whitespace-pre-wrap">
                    {tool.result_preview}
                  </pre>
                </div>
              )}
            </div>
          );
        })}
      </div>
    );
  };

  return (
    <div className={`flex flex-col bg-dark-800 rounded-xl border border-dark-600 overflow-hidden ${
      isExpanded ? 'fixed inset-4 z-50' : 'h-full'
    }`}>
      {/* Header */}
      <div className="px-4 py-3 border-b border-dark-600 flex items-center justify-between bg-dark-800">
        <div className="flex items-center gap-2">
          <div className="w-8 h-8 bg-gradient-to-br from-primary-500 to-primary-700 rounded-lg flex items-center justify-center">
            <Bot className="w-4 h-4 text-white" />
          </div>
          <div>
            <span className="font-medium text-white">Claude Code</span>
            {selectedStudy && (
              <span className="ml-2 text-xs bg-dark-700 px-2 py-0.5 rounded text-dark-300">
                {selectedStudy.id}
              </span>
            )}
          </div>
        </div>
        <div className="flex items-center gap-1">
          {messages.length > 0 && (
            <button
              onClick={clearConversation}
              className="p-2 text-dark-400 hover:text-white hover:bg-dark-700 rounded-lg transition-colors"
              title="Clear conversation"
            >
              <Trash2 className="w-4 h-4" />
            </button>
          )}
          {onToggleExpand && (
            <button
              onClick={onToggleExpand}
              className="p-2 text-dark-400 hover:text-white hover:bg-dark-700 rounded-lg transition-colors"
            >
              {isExpanded ? <Minimize2 className="w-4 h-4" /> : <Maximize2 className="w-4 h-4" />}
            </button>
          )}
          {onClose && (
            <button
              onClick={onClose}
              className="p-2 text-dark-400 hover:text-white hover:bg-dark-700 rounded-lg transition-colors"
            >
              <X className="w-4 h-4" />
            </button>
          )}
        </div>
      </div>

      {/* API Status Warning */}
      {apiAvailable === false && (
        <div className="px-4 py-3 bg-yellow-900/20 border-b border-yellow-800/30 flex items-center gap-2">
          <AlertCircle className="w-4 h-4 text-yellow-500" />
          <span className="text-yellow-400 text-sm">
            {error || 'Claude API not available'}
          </span>
        </div>
      )}

      {/* Messages */}
      <div className="flex-1 overflow-y-auto p-4 space-y-4">
        {messages.length === 0 ? (
          <div className="text-center py-8">
            <Sparkles className="w-12 h-12 mx-auto mb-4 text-primary-400 opacity-50" />
            <p className="text-dark-300 mb-2">Ask me anything about your optimization</p>
            <p className="text-dark-500 text-sm mb-6">
              I can analyze results, explain concepts, and help you improve your designs.
            </p>

            {/* Suggestions */}
            {suggestions.length > 0 && (
              <div className="flex flex-wrap gap-2 justify-center max-w-md mx-auto">
                {suggestions.slice(0, 6).map((suggestion, i) => (
                  <button
                    key={i}
                    onClick={() => sendMessage(suggestion)}
                    disabled={isLoading || apiAvailable === false}
                    className="px-3 py-1.5 bg-dark-700 hover:bg-dark-600 disabled:opacity-50
                             rounded-lg text-sm text-dark-300 hover:text-white transition-colors
                             border border-dark-600 hover:border-dark-500"
                  >
                    {suggestion}
                  </button>
                ))}
              </div>
            )}
          </div>
        ) : (
          messages.map((msg) => (
            <div
              key={msg.id}
              className={`flex gap-3 ${msg.role === 'user' ? 'justify-end' : ''}`}
            >
              {msg.role === 'assistant' && (
                <div className="w-8 h-8 rounded-lg bg-gradient-to-br from-primary-500 to-primary-700 flex items-center justify-center flex-shrink-0">
                  <Bot className="w-4 h-4 text-white" />
                </div>
              )}
              <div
                className={`max-w-[85%] rounded-lg ${
                  msg.role === 'user'
                    ? 'bg-primary-600 text-white px-4 py-2'
                    : 'bg-dark-700 text-dark-200 px-4 py-3'
                }`}
              >
                {msg.role === 'assistant' ? (
                  <>
                    <div className="prose prose-sm prose-invert max-w-none">
                      <ReactMarkdown
                        remarkPlugins={[remarkGfm]}
                        components={{
                          // Simplified markdown styling for chat
                          p: ({ children }) => <p className="mb-2 last:mb-0">{children}</p>,
                          ul: ({ children }) => <ul className="list-disc list-inside mb-2">{children}</ul>,
                          ol: ({ children }) => <ol className="list-decimal list-inside mb-2">{children}</ol>,
                          li: ({ children }) => <li className="mb-1">{children}</li>,
                          code: ({ inline, children }: any) =>
                            inline ? (
                              <code className="px-1 py-0.5 bg-dark-600 rounded text-primary-400 text-xs">
                                {children}
                              </code>
                            ) : (
                              <pre className="p-2 bg-dark-800 rounded overflow-x-auto text-xs my-2">
                                <code>{children}</code>
                              </pre>
                            ),
                          table: ({ children }) => (
                            <div className="overflow-x-auto my-2">
                              <table className="text-xs border-collapse">{children}</table>
                            </div>
                          ),
                          th: ({ children }) => (
                            <th className="border border-dark-600 px-2 py-1 bg-dark-800 text-left">
                              {children}
                            </th>
                          ),
                          td: ({ children }) => (
                            <td className="border border-dark-600 px-2 py-1">{children}</td>
                          ),
                          strong: ({ children }) => <strong className="text-white">{children}</strong>,
                          h1: ({ children }) => <h1 className="text-lg font-bold text-white mt-4 mb-2">{children}</h1>,
                          h2: ({ children }) => <h2 className="text-base font-semibold text-white mt-3 mb-2">{children}</h2>,
                          h3: ({ children }) => <h3 className="text-sm font-semibold text-white mt-2 mb-1">{children}</h3>,
                        }}
                      >
                        {msg.content}
                      </ReactMarkdown>
                    </div>
                    {renderToolCalls(msg.toolCalls, msg.id)}
                  </>
                ) : (
                  <p>{msg.content}</p>
                )}
              </div>
              {msg.role === 'user' && (
                <div className="w-8 h-8 rounded-lg bg-dark-600 flex items-center justify-center flex-shrink-0">
                  <User className="w-4 h-4 text-dark-300" />
                </div>
              )}
            </div>
          ))
        )}

        {/* Loading indicator */}
        {isLoading && (
          <div className="flex gap-3">
            <div className="w-8 h-8 rounded-lg bg-gradient-to-br from-primary-500 to-primary-700 flex items-center justify-center">
              <Loader2 className="w-4 h-4 text-white animate-spin" />
            </div>
            <div className="bg-dark-700 rounded-lg px-4 py-3 text-dark-400">
              <span className="inline-flex items-center gap-2">
                <span>Thinking</span>
                <span className="flex gap-1">
                  <span className="w-1.5 h-1.5 bg-dark-400 rounded-full animate-bounce" style={{ animationDelay: '0ms' }} />
                  <span className="w-1.5 h-1.5 bg-dark-400 rounded-full animate-bounce" style={{ animationDelay: '150ms' }} />
                  <span className="w-1.5 h-1.5 bg-dark-400 rounded-full animate-bounce" style={{ animationDelay: '300ms' }} />
                </span>
              </span>
            </div>
          </div>
        )}

        {/* Error message */}
        {error && !isLoading && (
          <div className="flex items-center gap-2 px-4 py-2 bg-red-900/20 border border-red-800/30 rounded-lg">
            <AlertCircle className="w-4 h-4 text-red-400" />
            <span className="text-red-400 text-sm">{error}</span>
          </div>
        )}

        <div ref={messagesEndRef} />
      </div>

      {/* Input */}
      <div className="p-4 border-t border-dark-600 bg-dark-800">
        <div className="flex gap-2">
          <input
            ref={inputRef}
            type="text"
            value={input}
            onChange={(e) => setInput(e.target.value)}
            onKeyDown={(e) => {
              if (e.key === 'Enter' && !e.shiftKey) {
                e.preventDefault();
                sendMessage();
              }
            }}
            placeholder={apiAvailable === false ? 'API not available...' : 'Ask about your optimization...'}
            disabled={isLoading || apiAvailable === false}
            className="flex-1 px-4 py-2.5 bg-dark-700 border border-dark-600 rounded-lg
                     text-white placeholder-dark-400 focus:outline-none focus:border-primary-500
                     disabled:opacity-50 disabled:cursor-not-allowed"
          />
          <button
            onClick={() => sendMessage()}
            disabled={!input.trim() || isLoading || apiAvailable === false}
            className="px-4 py-2.5 bg-primary-600 hover:bg-primary-500 disabled:opacity-50
                     disabled:cursor-not-allowed text-white rounded-lg transition-colors"
          >
            {isLoading ? (
              <Loader2 className="w-5 h-5 animate-spin" />
            ) : (
              <Send className="w-5 h-5" />
            )}
          </button>
        </div>
        <p className="text-xs text-dark-500 mt-2 text-center">
          Claude can query your study data, analyze results, and help improve your optimization.
        </p>
      </div>
    </div>
  );
};

export default ClaudeChat;