// Message component for individual messages
const Message = ({ message, isLoading, isLastAssistant }) => {
  const isUser = message.role === 'user';
  const [urlMap, setUrlMap] = React.useState({});
  
  return (
    <div className={`my-4 ${isUser ? 'ml-auto' : 'mr-auto'} max-w-[85%] sm:max-w-[80%]`}>
      <div className={`p-4 rounded-2xl shadow-sm ${
        isUser 
          ? 'bg-gradient-to-br from-blue-500/90 to-indigo-600/90 text-white' 
          : 'bg-white dark:bg-gray-800 text-gray-800 dark:text-gray-100 border border-gray-200 dark:border-gray-700'
      }`}>
        {/* Render user messages as plain text, assistant as Markdown using marked */}
        {isUser ? (
          <div className="whitespace-pre-wrap break-words leading-relaxed">{message.content}</div>
        ) : (
          <>
              <div
                className="prose prose-sm dark:prose-invert max-w-none"
                dangerouslySetInnerHTML={{ __html: window.marked ? window.marked.parse(message.content) : message.content }}
              />
              {((isLoading && isLastAssistant) || (message.content && isLoading)) && <TypingIndicator />}

          </>
        )}
        
        {/* Sources */}
        {message.sources && message.sources.length > 0 && (
          <div className="mt-3 pt-3 border-t border-white/20 dark:border-gray-700 text-xs opacity-90">
            <div className="font-semibold mb-1">Sources:</div>
            <ul className="space-y-1.5">
              {message.sources.map((source, i) => {
                let sourceText = '';
                let link = null;
                if (source.source_type === 'hydrocarbon_article') {
                  sourceText = `${source.title || 'Article'}`;
                  if (source.content_id && urlMap[source.content_id]) {
                    link = urlMap[source.content_id];
                  } else if (source.url) {
                    link = source.url;
                  } else if (source.source) {
                    link = source.source;
                  }
                } else if (source.source_type === 'PDF') {
                  sourceText = `${source.source || 'Document'}`;
                  if (source.page > 0) sourceText += ` (page ${source.page})`;
                } else {
                  sourceText = source.source || 'Unknown source';
                }
                return (
                  <li key={i} className="flex items-start">
                    <span className="inline-block mr-1">•</span>
                    <span>
                      {link ? (
                        <a href={link} target="_blank" rel="noopener noreferrer" className="text-blue-600 underline hover:text-blue-800">
                          {sourceText}
                        </a>
                      ) : (
                        sourceText
                      )}
                    </span>
                  </li>
                );
              })}
            </ul>
          </div>
        )}
      </div>
    </div>
  );
};

// Typing indicator component
const TypingIndicator = () => (
  <div className="flex items-center space-x-2 p-4">
    <div className="w-2 h-2 rounded-full bg-blue-400 animate-pulse" style={{animationDelay: '0ms'}}></div>
    <div className="w-2 h-2 rounded-full bg-blue-400 animate-pulse" style={{animationDelay: '300ms'}}></div>
    <div className="w-2 h-2 rounded-full bg-blue-400 animate-pulse" style={{animationDelay: '600ms'}}></div>
  </div>
);

// Main App component
function App() {
  const [messages, setMessages] = React.useState([]);
  const [input, setInput] = React.useState('');
  const [isLoading, setIsLoading] = React.useState(false);
  const [threadId, setThreadId] = React.useState(localStorage.getItem('threadId') || '');
  
  const messagesEndRef = React.useRef(null);
  const inputRef = React.useRef(null);
  
  // Focus input on load
  React.useEffect(() => {
    if (inputRef.current) {
      inputRef.current.focus();
    }
  }, []);
  
  // Scroll to bottom when messages change
  React.useEffect(() => {
    if (messagesEndRef.current) {
      messagesEndRef.current.scrollIntoView({ behavior: 'smooth' });
    }
  }, [messages]);
  
  // Add custom CSS
  React.useEffect(() => {
    const style = document.createElement('style');
    style.textContent = `
      @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
      
      body {
        font-family: 'Inter', system-ui, -apple-system, sans-serif;
        -webkit-font-smoothing: antialiased;
      }
      
      .btn-shine {
        position: relative;
        overflow: hidden;
      }
      
      .btn-shine::after {
        content: '';
        position: absolute;
        top: -50%;
        left: -50%;
        width: 200%;
        height: 200%;
        background: linear-gradient(
          to bottom right,
          rgba(255, 255, 255, 0) 0%,
          rgba(255, 255, 255, 0.1) 100%
        );
        transform: rotate(30deg);
        transition: transform 0.5s;
      }
      
      .btn-shine:hover::after {
        transform: rotate(30deg) translate(10%, 10%);
      }
      
      .animate-pulse {
        animation: pulse 1.5s cubic-bezier(0.4, 0, 0.6, 1) infinite;
      }
      
      @keyframes pulse {
        0%, 100% {
          opacity: 1;
        }
        50% {
          opacity: 0.3;
        }
      }
    `;
    document.head.appendChild(style);
    
    return () => {
      document.head.removeChild(style);
    };
  }, []);
  
  const handleSend = async () => {
    if (!input.trim() || isLoading) return;
    
    const userMessage = input;
    setInput('');
    setIsLoading(true);
    
    // Add user message to state
    setMessages(prevMessages => [...prevMessages, { role: 'user', content: userMessage }]);
    
    try {
      // Call API
      const response = await fetch('/chat', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          user_input: userMessage,
          thread_id: threadId || null
        }),
      });
      
      if (!response.ok) {
        const errorText = await response.text();
        console.error(`API error (${response.status}):`, errorText);
        throw new Error(`API error: ${response.status} - ${errorText}`);
      }
      
      const data = await response.json();
      
      // Log entire response for debugging
      console.log("API Response:", data);
      
      // Validate response structure
      if (!data.response) {
        console.error("Invalid API response - missing 'response' field:", data);
        throw new Error("Invalid API response format");
      }
      
      // Always use the thread ID from the server
      if (data.thread_id) {
        console.log(`Using thread ID: ${data.thread_id}`);
        setThreadId(data.thread_id);
        localStorage.setItem('threadId', data.thread_id);
      }
      
      // Add bot response to state, ensuring sources are properly attached
      setMessages(prevMessages => [
        ...prevMessages, 
        { 
          role: 'assistant', 
          content: data.response,
          sources: Array.isArray(data.sources) ? data.sources : [] // Ensure sources is an array
        }
      ]);
    } catch (error) {
      console.error('Error during API call:', error);
      setMessages(prevMessages => [
        ...prevMessages, 
        { role: 'system', content: `Error: ${error.message || 'An unknown error occurred'}` }
      ]);
    } finally {
      setIsLoading(false);
      // Focus input after response
      if (inputRef.current) {
        inputRef.current.focus();
      }
    }
  };
  
  // New streaming version of handleSend that uses the /stream_chat endpoint
  const handleStreamSend = async () => {
    if (!input.trim() || isLoading) return;
    
    const userMessage = input;
    setInput('');
    setIsLoading(true);
    
    // Add user message to state
    setMessages(prevMessages => [...prevMessages, { role: 'user', content: userMessage }]);
    // Always add the assistant message bubble immediately
    setMessages(prevMessages => [
      ...prevMessages,
      { role: 'assistant', content: '', sources: [] }
    ]);
    
    let firstTokenReceived = false;
    try {
      // Make the request to the streaming endpoint
      const response = await fetch('/stream_chat', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          user_input: userMessage,
          thread_id: threadId || null
        }),
      });
      
      if (!response.ok) {
        const errorText = await response.text();
        console.error(`API error (${response.status}):`, errorText);
        throw new Error(`API error: ${response.status} - ${errorText}`);
      }
      
      // Get the reader from the response body stream
      const reader = response.body.getReader();
      const decoder = new TextDecoder();
      let buffer = '';
      
      // Process the stream
      while (true) {
        const { value, done } = await reader.read();
        if (done) break;
        
        // Decode the chunk and add it to our buffer
        buffer += decoder.decode(value, { stream: true });
        
        // Process all complete events in the buffer
        const events = buffer.split('\n\n');
        buffer = events.pop() || ''; // Keep the last incomplete chunk in the buffer
        for (const eventStr of events) {
          if (!eventStr.trim()) continue;
          // Only look for data: lines
          const dataLine = eventStr.split('\n').find(line => line.startsWith('data:'));
          if (!dataLine) continue;
          let data;
          try {
            data = JSON.parse(dataLine.slice(5).trim());
          } catch (e) {
            continue;
          }
          // Update assistant message as soon as a token arrives
          if (data.token !== undefined) {
            setMessages(prevMessages => {
              const newMessages = [...prevMessages];
              const assistantMessageIndex = newMessages.length - 1;
              if (assistantMessageIndex >= 0 && newMessages[assistantMessageIndex].role === 'assistant') {
                newMessages[assistantMessageIndex].content += data.token;
              }
              return newMessages;
            });
            if (!firstTokenReceived && data.token && data.token !== "") {
              // setIsLoading(false);
              firstTokenReceived = true;
            }
          }
          // Attach sources when received
          if (data.sources) {
            setMessages(prevMessages => {
              const newMessages = [...prevMessages];
              const assistantMessageIndex = newMessages.length - 1;
              if (assistantMessageIndex >= 0 && newMessages[assistantMessageIndex].role === 'assistant') {
                newMessages[assistantMessageIndex].sources = data.sources;
              }
              return newMessages;
            });
          }
          if (data.error) {
            setIsLoading(false);
          }
        }
      }
    } catch (error) {
      console.error('Error during streaming API call:', error);
      // Update the last message to show the error, or add a new system message
      setMessages(prevMessages => {
        const newMessages = [...prevMessages];
        const lastMessageIndex = newMessages.length - 1;
        
        // If the last message is an empty assistant message, replace it
        if (lastMessageIndex >= 0 && 
            newMessages[lastMessageIndex].role === 'assistant' && 
            !newMessages[lastMessageIndex].content.trim()) {
          newMessages.pop(); // Remove the empty assistant message
        }
        
        // Add error message
        newMessages.push({ 
          role: 'system', 
          content: `Error: ${error.message || 'An unknown error occurred'}` 
        });
        
        return newMessages;
      });
    } finally {
      setIsLoading(false);
      // Focus input after response
      if (inputRef.current) {
        inputRef.current.focus();
      }
    }
  };
  
  const handleNewChat = () => {
    if (confirm('Start a new chat? This will clear the current conversation.')) {
      setMessages([]);
      setThreadId('');
      localStorage.removeItem('threadId');
    }
  };
  
  // Use the streaming version by default
  const handleSendMessage = handleStreamSend;
  
  return (
    <div className="h-screen flex flex-col bg-gradient-to-br from-gray-50 to-gray-100 dark:from-gray-900 dark:to-black">
      {/* Header - fixed height */}
      <div className="h-16 flex-none bg-white dark:bg-gray-900 border-b border-gray-200 dark:border-gray-800 flex items-center justify-between px-6 shadow-sm z-20">
        <div className="flex items-center space-x-2">
          <div className="w-2 h-2 rounded-full bg-blue-500 animate-pulse"></div>
          <h1 className="text-xl font-bold text-gray-900 dark:text-white tracking-tight">AI Assistant</h1>
        </div>
        <button 
          onClick={handleNewChat}
          className="px-4 py-2 bg-gradient-to-r from-blue-500 to-indigo-500 hover:from-blue-600 hover:to-indigo-600 rounded-full text-sm text-white font-medium shadow-sm transition-all duration-200 btn-shine"
        >
          New Chat
        </button>
      </div>
      
      {/* Messages container - takes remaining space, scrollable */}
      <div className="relative flex-1">
        <div className="absolute inset-0 overflow-y-auto pb-40 scrollbar-thin scrollbar-thumb-gray-300 dark:scrollbar-thumb-gray-700">
          <div className="max-w-3xl mx-auto px-4 py-8">
            {messages.length === 0 ? (
              <div className="text-center my-20">
                <div className="w-20 h-20 mx-auto rounded-full bg-blue-100 dark:bg-blue-900/30 flex items-center justify-center mb-6">
                  <svg className="w-10 h-10 text-blue-600 dark:text-blue-400" fill="none" viewBox="0 0 24 24" stroke="currentColor">
                    <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M8 10h.01M12 10h.01M16 10h.01M9 16H5a2 2 0 01-2-2V6a2 2 0 012-2h14a2 2 0 012 2v8a2 2 0 01-2 2h-5l-5 5v-5z" />
                  </svg>
                </div>
                <h3 className="text-xl font-semibold text-gray-900 dark:text-white mb-2">How can I help you today?</h3>
                <p className="text-gray-600 dark:text-gray-400 max-w-md mx-auto">Ask me anything about the documents I've processed, and I'll provide insights based on the available information.</p>
              </div>
            ) : (
              <>
                {messages.map((msg, index) => (
                  <Message
                    key={index}
                    message={msg}
                    isLoading={isLoading && index === messages.length - 1 && msg.role === 'assistant'}
                    isLastAssistant={index === messages.length - 1 && msg.role === 'assistant'}
                  />
                ))}
                <div ref={messagesEndRef} className="h-px" />
              </>
            )}
          </div>
        </div>
      </div>
      
      {/* Input container - fixed at bottom with set height and solid background */}
      <div className="absolute bottom-0 left-0 right-0 pointer-events-none z-10">
        <div className="bg-gray-50 dark:bg-gray-900 w-full py-4 pointer-events-auto">
          <div className="max-w-3xl mx-auto px-4 pb-2 pt-2">
            <div className="bg-white dark:bg-gray-800 rounded-2xl shadow-lg border border-gray-200 dark:border-gray-700 overflow-hidden">
              <div className="flex">
                <textarea
                  ref={inputRef}
                  value={input}
                  onChange={(e) => setInput(e.target.value)}
                  onKeyDown={(e) => {
                    if (e.key === 'Enter' && !e.shiftKey) {
                      e.preventDefault();
                      handleSendMessage();
                    }
                  }}
                  className="flex-grow p-4 bg-transparent outline-none resize-none text-gray-800 dark:text-white placeholder-gray-500 dark:placeholder-gray-400"
                  placeholder="Type your message here..."
                  disabled={isLoading}
                  style={{height: "80px", minHeight: "80px", maxHeight: "80px"}}
                />
                <div className="flex items-center pr-3">
                  <button 
                    onClick={handleSendMessage}
                    disabled={isLoading || !input.trim()}
                    className={`w-12 h-12 rounded-full flex items-center justify-center transition-all duration-200 ${
                      isLoading || !input.trim() 
                        ? 'bg-gray-200 dark:bg-gray-700 text-gray-400 cursor-not-allowed' 
                        : 'bg-gradient-to-r from-blue-500 to-indigo-500 text-white shadow-md hover:shadow-lg btn-shine'
                    }`}
                  >
                    <svg className="w-5 h-5" viewBox="0 0 20 20" fill="currentColor">
                      <path d="M10.894 2.553a1 1 0 00-1.788 0l-7 14a1 1 0 001.169 1.409l5-1.429A1 1 0 009 15.571V11a1 1 0 112 0v4.571a1 1 0 00.725.962l5 1.428a1 1 0 001.17-1.408l-7-14z" />
                    </svg>
                  </button>
                </div>
              </div>
            </div>
          </div>
        </div>
      </div>
    </div>
  );
}

// Render the App
const rootElement = document.getElementById('root');
const root = ReactDOM.createRoot(rootElement);
root.render(<App />);