import React, { useState, useEffect, useRef } from 'react';
import { useAuth } from '../services/AuthContext';
import { useNavigate } from 'react-router-dom';

const processMessages = (messageObject) => {
  return (previousMessages) => {
    if (previousMessages.length > 0 && previousMessages[previousMessages.length - 1].sender === messageObject.sender) {
      const newMessage = {text:previousMessages[previousMessages.length - 1].text+ ' ' + messageObject.text, sender: messageObject.sender};
      return [...previousMessages.slice(0, -1), newMessage];
    } else {
      return [...previousMessages, messageObject];
    }
  }
}

const Chat = () => {
  const { token, logout } = useAuth();
  const navigate = useNavigate();
  const [messages, setMessages] = useState([]);
  const [input, setInput] = useState('');
  const [isConnected, setIsConnected] = useState(false);
  const [connectionAttempts, setConnectionAttempts] = useState(0);
  const [isReconnecting, setIsReconnecting] = useState(false);
  const websocket = useRef(null);
  const maxReconnectAttempts = 5;

  const connectWebSocket = () => {
    if (websocket.current?.readyState === WebSocket.OPEN) {
      console.log('WebSocket is already connected');
      return;
    }

    console.log('Attempting to connect WebSocket...');
    const ws = new WebSocket('ws://localhost:8000/ws');
    websocket.current = ws;

    ws.onopen = async () => {
      console.log('WebSocket connected, sending auth token...');
      // Send authentication token
      ws.send(JSON.stringify({
        type: 'authentication',
        token: token
      }));
    };

    ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        console.log('Received message:', data);

        switch (data.type) {
          case 'connection_established':
            setIsConnected(true);
            setConnectionAttempts(0);
            setIsReconnecting(false);
            break;
          case 'message':
            setMessages(processMessages({
              text:data.message,
              sender: data.sender
            }));
            break;
          case 'error':
            console.error('Server error:', data.message);
            if (data.message.includes('Authentication failed')) {
              logout();
              navigate('/login');
            }
            break;
          default:
            console.log('Unhandled message type:', data.type);
        }
      } catch (error) {
        console.error('Error processing message:', error);
        setMessages(processMessages({
          text: event.data,
          sender: 'ai'
        }));
      }
    };

    ws.onclose = (event) => {
      console.log('WebSocket closed:', event);
      setIsConnected(false);

      if (event.code === 1008) {
        // Authentication failed
        console.error('WebSocket authentication failed');
        logout();
        navigate('/login');
      } else if (connectionAttempts < maxReconnectAttempts) {

        // Normal closure or error, attempt to reconnect
        setIsReconnecting(true);
        setConnectionAttempts(prev => prev + 1);
        setTimeout(connectWebSocket, 2000 * Math.min(connectionAttempts + 1, 5));
      }
    };

    ws.onerror = (error) => {
      console.error('WebSocket error:', error);
    };
  };

  useEffect(() => {
    connectWebSocket();

    return () => {
      if (websocket.current) {
        websocket.current.close();
      }
    };
  }, [token]);

  const sendMessage = (e) => {
    e.preventDefault();
    if (!input.trim() || !isConnected) return;

    const message = {
      type: 'message',
      content: input
    };

    setMessages(processMessages({
      text: input,
      sender: 'user'
    }));
    websocket.current.send(JSON.stringify(message));
    setInput('');
  };
  console.log('MESSAGES****', messages);
  // Rest of your component (file upload handler, UI rendering, etc.)
  return (
    <div className="flex flex-col h-screen bg-gray-100">
      

      {/* Messages area */}
      <div className="flex-1 overflow-y-auto p-4 space-y-4">
      {isReconnecting && <span className="text-yellow-500">Reconnecting...</span>}
        {(messages || []).map((message, index) => (
          <div
            key={index}
            className={`p-3 rounded-lg max-w-[80%] ${message.sender === 'user'
                ? 'ml-auto bg-blue-500 text-white'
                : message.sender === 'ai'
                  ? 'bg-gray-200'
                  : 'bg-yellow-100 mx-auto'
              }`}
          >
            {message.text}
          </div>
        ))}
      </div>

      {/* Input area */}
      <div className="p-4 bg-white border-t">
        <form onSubmit={sendMessage} className="flex gap-4">
          <input
            type="text"
            value={input}
            onChange={(e) => setInput(e.target.value)}
            placeholder="Type your message..."
            className="flex-1 p-2 border rounded"
            disabled={!isConnected}
          />
          <button
            type="submit"
            className="px-4 py-2 bg-blue-500 text-white rounded hover:bg-blue-600 disabled:bg-gray-400"
            disabled={!isConnected}
          >
            Send
          </button>

          <button
            onClick={() => {
              logout();
              navigate('/login');
            }}
            className="px-4 py-2 text-white bg-red-500 rounded hover:bg-red-600"
          >
            Logout
          </button>
          <span className={`h-3 w-3 rounded-full ${isConnected ? 'bg-green-500' : 'bg-red-500'}`}></span>
        </form>
      </div>
    </div>
  );
};

export default Chat;