frontend/src/components/InferencePanel.tsx

// src/components/InferencePanel.tsx
import React, { useState, useEffect, useRef } from 'react';
import axios from 'axios';

const API_URL = import.meta.env.VITE_API_URL || 'http://localhost:8000';

// Define message types
interface Message {
  id: string;
  sender: 'user' | 'ai';
  text: string;
  domain?: string;
  confidence?: number;
  modelUsed?: string;
  sourceType?: string;
  thinkingProcess?: string;
  thinkingSteps?: string[];
  isEditing?: boolean;
  editedText?: string;
}

interface InferencePanelProps {
  activeDomain: string;
  currentSessionId?: string;
}

const InferencePanel: React.FC<InferencePanelProps> = ({ activeDomain }) => {
  const [sessionId] = useState(`ws-session-${Date.now()}-${Math.random()}`);
  const [socket, setSocket] = useState<WebSocket | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [messages, setMessages] = useState<Message[]>([]);
  const [inputValue, setInputValue] = useState('');
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [ragMode, setRagMode] = useState<'direct' | 'rag'>('rag');
  const [currentThinkingSteps, setCurrentThinkingSteps] = useState<string[]>([]);
  const [availableDomains, setAvailableDomains] = useState<any[]>([]);
  const [editingMessageId, setEditingMessageId] = useState<string | null>(null);
  const [editingText, setEditingText] = useState<string>('');

  const messagesEndRef = useRef<null | HTMLDivElement>(null);

  // Load available domains
  useEffect(() => {
    const fetchDomains = async () => {
      try {
        const response = await axios.get(`${API_URL}/api/config/domains`);
        setAvailableDomains(response.data);
      } catch (error) {
        console.error('Failed to fetch domains:', error);
      }
    };
    fetchDomains();
  }, []);

  // Effect to manage WebSocket connection
  useEffect(() => {
    const wsUrl = `ws://localhost:8000/api/questions/ws/${sessionId}`;
    const ws = new WebSocket(wsUrl);

    ws.onopen = () => {
      console.log('WebSocket connected');
      setIsConnected(true);
      setError(null);
    };

    ws.onmessage = (event) => {
      const data = JSON.parse(event.data);
      handleSocketMessage(data);
    };

    ws.onclose = () => {
      console.log('WebSocket disconnected');
      setIsConnected(false);
    };

    ws.onerror = (error) => {
      console.error('WebSocket error:', error);
      setIsConnected(false);
      setIsLoading(false);
      setError('Could not connect to the inference server.');
    };

    setSocket(ws);

    // Cleanup on unmount
    return () => {
      ws.close();
    };
  }, [sessionId]);

  // Effect to scroll to the bottom of messages
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages, currentThinkingSteps]);


  const handleSocketMessage = (data: any) => {
    switch (data.type) {
      case 'processing':
        setCurrentThinkingSteps(prev => [...prev, 'Processing...']);
        break;
      
      case 'thinking':
        // Update current thinking steps for real-time display
        setCurrentThinkingSteps(prev => [...prev, data.step]);
        break;

      case 'token':
        setMessages(prev => {
          const lastMessage = prev[prev.length - 1];
          if (lastMessage && lastMessage.sender === 'ai' && !lastMessage.sourceType) { // Ensure it's the current AI response
            // Append token to the last AI message
            const updatedMessages = [...prev];
            updatedMessages[prev.length - 1] = { ...lastMessage, text: lastMessage.text + data.content };
            return updatedMessages;
          } else {
            // Start a new AI message (should be rare, but good for robustness)
            return [...prev, {
              id: `ai-${Date.now()}-${Math.random()}`,
              sender: 'ai',
              text: data.content || '',
              domain: activeDomain
            }];
          }
        });
        break;
      
      case 'meta': // Metadata sent before or during streaming
        setMessages(prev => {
          const lastMessage = prev[prev.length - 1];
          if (lastMessage && lastMessage.sender === 'ai') {
            const updatedMessages = [...prev];
            updatedMessages[prev.length - 1] = { 
              ...lastMessage, 
              confidence: data.confidence,
              modelUsed: data.model_used,
              sourceType: data.source_type,
              thinkingProcess: data.thinking_process // Changed from data.thinking
            };
            return updatedMessages;
          }
          return prev;
        });
        break;

      case 'response': // Final complete answer
        setIsLoading(false);
        setCurrentThinkingSteps([]); // Clear real-time thinking steps
        setMessages(prev => {
          const lastMessage = prev[prev.length - 1];
          if (lastMessage && lastMessage.sender === 'ai') {
             const updatedMessages = [...prev];
             // Ensure final metadata is set
             updatedMessages[prev.length - 1] = {
               ...lastMessage,
               text: data.response, // Final complete response
               confidence: data.confidence || lastMessage.confidence,
               modelUsed: data.model_used || lastMessage.modelUsed,
               sourceType: data.source_type || lastMessage.sourceType,
               thinkingProcess: data.thinking_process || lastMessage.thinkingProcess, // Changed from data.thinking
               thinkingSteps: currentThinkingSteps // Attach all thinking steps to the final message
             };
             return updatedMessages;
          }
          // If no streaming tokens received, create a new message
          return [...prev, {
            id: `ai-${Date.now()}-${Math.random()}`,
            sender: 'ai',
            text: data.response,
            domain: activeDomain,
            confidence: data.confidence,
            modelUsed: data.model_used,
            sourceType: data.source_type,
            thinkingProcess: data.thinking_process, // Changed from data.thinking
            thinkingSteps: currentThinkingSteps
          }];
        });
        break;

      case 'error':
        setIsLoading(false);
        setCurrentThinkingSteps([]);
        setError(`WebSocket Error: ${data.error}`); // Set error state for prominent display
        break;
    }
  };

  const handleDomainChange = (messageId: string, newDomain: string) => {
    setMessages(prev => prev.map(msg =>
      msg.id === messageId ? { ...msg, domain: newDomain } : msg
    ));
  };

  const handleStartEdit = (msg: Message) => {
    setEditingMessageId(msg.id);
    setEditingText(msg.text);
  };

  const handleSaveEdit = (msgId: string) => {
    setMessages(prev => prev.map(msg =>
      msg.id === msgId ? { ...msg, text: editingText } : msg
    ));
    setEditingMessageId(null);
  };

  const handleCancelEdit = () => {
    setEditingMessageId(null);
    setEditingText('');
  };

  const handleSendMessage = () => {
    if (inputValue.trim() && socket && isConnected && !isLoading) {
      // Add user message to chat
      const newMessage: Message = {
        id: `msg-${Date.now()}-${Math.random()}`,
        sender: 'user',
        text: inputValue,
        domain: activeDomain
      };
      setMessages(prev => [...prev, newMessage]);

      // Send question to WebSocket server
      socket.send(JSON.stringify({
        type: 'question',
        question: inputValue,
        stream: true, // Always stream
        domain_id: activeDomain, // Use the active domain from props
        rag_mode: ragMode // Add the selected RAG mode
      }));

      setIsLoading(true);
      setInputValue('');
      setCurrentThinkingSteps([]); // Clear previous thinking steps display
    }
  };

  const getSourceTypeColor = (sourceType?: string) => {
    switch (sourceType) {
      case 'db_augmented': return 'bg-blue-600';
      case 'ai_internal_weights': return 'bg-purple-600';
      case 'web_search': return 'bg-yellow-600';
      default: return 'bg-gray-600';
    }
  };

  return (
    <div className="chat-container">
      <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center', paddingBottom: '0.5rem' }}>
        <h3 style={{ margin: 0, color: '#00aaff' }}>Inference</h3>
        <span
          className={`connection-status ${isConnected ? 'connected' : 'disconnected'}`}
          style={{ fontSize: '0.75rem', padding: '0.2rem 0.5rem', borderRadius: '10px' }}
        >
          {isConnected ? '● Connected' : '○ Disconnected'}
        </span>
      </div>
      {error && (
        <div style={{ color: 'red', marginBottom: '1rem', padding: '0.5rem', backgroundColor: 'rgba(255,0,0,0.1)', borderRadius: '5px' }}>
          {error}
        </div>
      )}
      <div className="messages">
        {messages.map((msg) => (
          <div key={msg.id} className={`message ${msg.sender}`}>
            {editingMessageId === msg.id ? (
              <div>
                <textarea
                  value={editingText}
                  onChange={(e) => setEditingText(e.target.value)}
                  style={{
                    width: '100%',
                    minHeight: '100px',
                    padding: '8px',
                    backgroundColor: '#2a2a2a',
                    border: '1px solid #444',
                    borderRadius: '4px',
                    color: '#fff',
                    fontSize: '14px'
                  }}
                />
                <div style={{ marginTop: '8px', display: 'flex', gap: '8px' }}>
                  <button onClick={() => handleSaveEdit(msg.id)} style={{ padding: '4px 8px', fontSize: '12px' }}>保存</button>
                  <button onClick={handleCancelEdit} style={{ padding: '4px 8px', fontSize: '12px' }}>キャンセル</button>
                </div>
              </div>
            ) : (
              <div>
                {msg.text}
                {msg.sender === 'ai' && (
                  <div style={{ marginTop: '8px', display: 'flex', gap: '8px', alignItems: 'center', flexWrap: 'wrap' }}>
                    {msg.domain && availableDomains.length > 0 && (
                      <>
                        <span style={{ fontSize: '12px', color: '#888' }}>Domain:</span>
                        <select
                          value={msg.domain}
                          onChange={(e) => handleDomainChange(msg.id, e.target.value)}
                          style={{
                            padding: '4px 8px',
                            fontSize: '12px',
                            backgroundColor: '#2a2a2a',
                            border: '1px solid #444',
                            borderRadius: '4px',
                            color: '#fff'
                          }}
                        >
                          {availableDomains.map(domain => (
                            <option key={domain.domain_id} value={domain.domain_id}>
                              {domain.name}
                            </option>
                          ))}
                        </select>
                      </>
                    )}
                    <button
                      onClick={() => handleStartEdit(msg)}
                      style={{
                        padding: '4px 8px',
                        fontSize: '12px',
                        backgroundColor: '#444',
                        border: 'none',
                        borderRadius: '4px',
                        color: '#fff',
                        cursor: 'pointer'
                      }}
                    >
                      編集
                    </button>
                  </div>
                )}
                {msg.sender === 'ai' && msg.confidence !== undefined && (
                  <div className="response-meta text-xs mt-2 pt-2 border-t border-gray-700 flex flex-wrap gap-2">
                    {msg.confidence !== undefined && (
                      <span className="confidence-badge bg-green-600">
                        Confidence: {(msg.confidence * 100).toFixed(1)}%
                      </span>
                    )}
                    {msg.modelUsed && (
                      <span className="memory-badge bg-gray-600">
                        Model: {msg.modelUsed}
                      </span>
                    )}
                    {msg.sourceType && (
                      <span className={`memory-badge ${getSourceTypeColor(msg.sourceType)}`}>
                        Source: {msg.sourceType.replace('_', ' ')}
                      </span>
                    )}
                  </div>
                )}
                {msg.sender === 'ai' && msg.thinkingSteps && msg.thinkingSteps.length > 0 && (
                  <div className="thinking-steps mt-2 pt-2 border-t border-gray-700">
                    <div className="thinking-header text-xs text-gray-400">Thinking Process:</div>
                    {msg.thinkingSteps.map((step, stepIndex) => (
                      <div key={stepIndex} className="thinking-step text-xs text-gray-500 pl-4 border-l-2 border-blue-500 my-1">
                        {step}
                      </div>
                    ))}
                  </div>
                )}
              </div>
            )}
          </div>
        ))}
        {/* Real-time thinking steps display */}
        {isLoading && currentThinkingSteps.length > 0 && (
          <div className="thinking-steps" style={{ background: '#2a2a2a', border: '1px solid #444', borderRadius: '6px', padding: '1rem', marginBottom: '1rem' }}>
            <div className="thinking-header" style={{ color: '#888', fontSize: '0.85rem', marginBottom: '0.5rem' }}>Thinking... <span className="loading-dots"><span>.</span><span>.</span><span>.</span></span></div>
            {currentThinkingSteps.map((step, i) => (
              <div key={i} className="thinking-step" style={{ color: '#aaa', fontSize: '0.85rem', padding: '0.25rem 0', paddingLeft: '1rem', borderLeft: '2px solid #007acc', marginBottom: '0.25rem' }}>{step}</div>
            ))}
          </div>
        )}
        <div ref={messagesEndRef} />
      </div>

      <div className="rag-mode-selector" style={{ display: 'flex', justifyContent: 'center', alignItems: 'center', marginBottom: '0.5rem', gap: '0.5rem' }}>
        <span style={{ fontSize: '0.8rem', color: '#ccc' }}>Query Mode:</span>
        <div style={{ display: 'flex', border: '1px solid #555', borderRadius: '15px', padding: '2px' }}>
          <button 
            onClick={() => setRagMode('direct')}
            className={`rag-mode-btn ${ragMode === 'direct' ? 'active' : ''}`}
            disabled={isLoading}
          >
            Direct DB
          </button>
          <button 
            onClick={() => setRagMode('rag')}
            className={`rag-mode-btn ${ragMode === 'rag' ? 'active' : ''}`}
            disabled={isLoading}
          >
            RAG Mode
          </button>
        </div>
      </div>

      <div className="input-area">
        <input
          type="text"
          placeholder="Ask NullAI anything..."
          value={inputValue}
          onChange={(e) => setInputValue(e.target.value)}
          onKeyPress={(e) => e.key === 'Enter' && handleSendMessage()}
          disabled={isLoading || !isConnected}
        />
        <button onClick={handleSendMessage} disabled={isLoading || !isConnected}>
          {!isConnected ? 'Connect.' : (isLoading ? 'Generating...' : 'Send')}
        </button>
      </div>
    </div>
  );
};

export default InferencePanel;