import React, { useEffect, createContext, useRef, useState, useContext, ReactNode } from 'react';
import { AccountContext } from './AccountContext';
import { AuthContext } from './AuthContext';

// ---
// Define Types and Interfaces
// ---

/**
 * Interface for the Account Context value.
 * Assumes account has at least an email property.
 */
interface IAccount {
  email: string | null;
  // Add other properties of the account object as needed
  // e.g., userId: string; name: string;
}

export interface ChatMessage {
  text: string;
  sender: 'user' | 'ai';
}

/**
 * Interface for the Auth Context value.
 */
interface IAuthContext {
  authenticated: boolean;
  loading: boolean;
}

/**
 * Type for the message data received from the WebSocket.
 * Adjust 'any' to a more specific type if the message structure is known.
 */
type WebSocketMessageData = string; // Assuming message.data is a string (e.g., JSON string)

/**
 * Type for a channel callback function.
 */
type ChannelCallback = (data: WebSocketMessageData) => void;

/**
 * Interface for the WebSocket Context value.
 * This defines the shape of the array returned by `useContext(WebSocketContext)`.
 */
interface IWebSocketContext {
  subscribe: (channel: string, callback: ChannelCallback) => void;
  unsubscribe: (channel: string) => void;
  socket: WebSocket | null;
  sendMessages: (messages: ChatMessage[]) => void;
}

/**
 * Props for the WebSocketProvider component.
 */
interface WebSocketProviderProps {
  children: ReactNode;
}

// ---
// Create the WebSocket Context
// ---

// Provide a default value that matches the IWebSocketContext interface.
// This is used when a component tries to consume the context without a provider.
const WebSocketContext = createContext<IWebSocketContext>({
  subscribe: () => {},
  unsubscribe: () => {},
  socket: null,
  sendMessages: () => {},
});

// ---
// WebSocket Provider Component
// ---

function WebSocketProvider({ children }: WebSocketProviderProps) {
  // Using useContext with explicit types for better type checking
  const { authenticated, loading } = useContext<IAuthContext>(AuthContext);
  const { account, setAccount } = useContext<{
    account: IAccount | null;
    setAccount: React.Dispatch<React.SetStateAction<IAccount | null>>;
  }>(AccountContext);

  // useRef for WebSocket instance
  const ws = useRef<WebSocket | null>(null);

  // useState for socket connection state
  const [socket, setSocket] = useState<WebSocket | null>(null);

  // useRef to store channel callbacks, mapping channel names to their callbacks
  const channels = useRef<{ [key: string]: ChannelCallback }>({});

  // useState for the current active channel (though not directly used in the current logic for message dispatch)
  const [currentChannel, setCurrentChannel] = useState<string>('');

  /**
   * Subscribes a callback function to a specific WebSocket channel.
   * @param channel The name of the channel to subscribe to.
   * @param callback The function to be called when a message is received on this channel.
   */
  const subscribe = (channel: string, callback: ChannelCallback) => {
    setCurrentChannel(channel); // This seems to track the last subscribed channel globally, not per-message.
    channels.current[channel] = callback;
  };

  /**
   * Unsubscribes a callback from a specific WebSocket channel.
   * @param channel The name of the channel to unsubscribe from.
   */
  const unsubscribe = (channel: string) => {
    delete channels.current[channel];
  };

  /**
   * Sends a message over the WebSocket connection.
   * Handles both text messages and file uploads (by converting file to base64).
   * @param message The text message to send.
   * @param account_iod The ID of the conversation this message belongs to.
   */
  const sendMessages = (messages: ChatMessage[]) => {
    if (socket && socket.readyState === WebSocket.OPEN) {
      const data = {
        messages: messages,
      };
      socket.send(JSON.stringify(data));
    } else {
      console.log('Error sending message. WebSocket is not open');
    }
  };

  // ---
  // WebSocket Initialization and Event Handling
  // ---
  useEffect(() => {
    if (account?.email) {
      // Ensure account and email exist before attempting to connect
      // Close any existing connection before creating a new one
      if (ws.current) {
        ws.current.close();
      }

      const wsUrl = new URL(import.meta.env.VITE_API_URL || 'ws://127.0.0.1:8010/ws/');
      wsUrl.protocol = wsUrl.protocol.replace('http', 'ws');

      ws.current = new WebSocket(
        `${wsUrl.origin}/ws/chat/${account.id}/?token=${localStorage.getItem('access_token')}`,
      );

      ws.current.onopen = () => {
        setSocket(ws.current);
        console.log('WebSocket connected');
      };

      ws.current.onclose = () => {
        setSocket(null);
        console.log('WebSocket disconnected');
      };

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

      ws.current.onmessage = (event: MessageEvent) => {
        const data: WebSocketMessageData = event.data;
        // The original logic assumes a single active chat channel is always present
        // and iterates Object.entries(channels.current)[0][0] to get its key.
        // This might be problematic if multiple channels are intended to be active
        // or if messages don't always correspond to the first subscribed channel.
        // For now, retaining the original logic, but consider refining based on backend message structure.
        const chatChannelKey = Object.keys(channels.current)[0];

        if (chatChannelKey && channels.current[chatChannelKey]) {
          // Call the callback associated with the identified channel
          channels.current[chatChannelKey](data);
        } else {
          console.log('No active chat channel subscribed or message format unexpected.');
          // Potentially handle generic messages or log the unhandled message
        }
      };

      // Cleanup function: Closes the WebSocket connection when the component unmounts
      // or when the 'account' dependency changes, triggering a re-run of this effect.
      return () => {
        if (ws.current) {
          ws.current.close();
          console.log('WebSocket cleaned up');
        }
      };
    } else if (!loading && !authenticated) {
      // If not authenticated and not loading, ensure WebSocket is closed if it was open.
      if (ws.current) {
        ws.current.close();
        ws.current = null;
        setSocket(null);
      }
    }
  }, [account, authenticated, loading]); // Dependencies: reconnect if account or auth status changes

  // ---
  // Render the Provider
  // ---
  return (
    <WebSocketContext.Provider value={{ subscribe, unsubscribe, socket, sendMessages }}>
      {children}
    </WebSocketContext.Provider>
  );
}

export { WebSocketContext, WebSocketProvider };
