import { create } from 'zustand';
import { persist } from 'zustand/middleware';

import { WebSocketClient, WebSocketState, WebSocketMessage, WebSocketMessageType } from '../features/chat/lib/websocket';

/**
 * EN: Message sender type
 * VI: Loại người gửi tin nhắn
 */
export enum MessageSender {
  /** EN: User message / VI: Tin nhắn từ người dùng */
  USER = 'user',
  /** EN: AI/Assistant message / VI: Tin nhắn từ AI/Assistant */
  ASSISTANT = 'assistant',
  /** EN: System message / VI: Tin nhắn hệ thống */
  SYSTEM = 'system',
}

/**
 * EN: Message status
 * VI: Trạng thái tin nhắn
 */
export enum MessageStatus {
  /** EN: Message is sending / VI: Tin nhắn đang được gửi */
  SENDING = 'sending',
  /** EN: Message is sent / VI: Tin nhắn đã được gửi */
  SENT = 'sent',
  /** EN: Message failed to send / VI: Gửi tin nhắn thất bại */
  FAILED = 'failed',
  /** EN: Message is delivered / VI: Tin nhắn đã được gửi đến */
  DELIVERED = 'delivered',
  /** EN: Message is read / VI: Tin nhắn đã được đọc */
  READ = 'read',
}

/**
 * EN: Chat message interface
 * VI: Interface tin nhắn chat
 */
export interface ChatMessage {
  /** EN: Unique message identifier / VI: Mã định danh tin nhắn duy nhất */
  id: string;
  /** EN: Conversation ID this message belongs to / VI: ID cuộc trò chuyện mà tin nhắn này thuộc về */
  conversationId: string;
  /** EN: Message content / VI: Nội dung tin nhắn */
  content: string;
  /** EN: Message sender / VI: Người gửi tin nhắn */
  sender: MessageSender;
  /** EN: Message status / VI: Trạng thái tin nhắn */
  status: MessageStatus;
  /** EN: Message creation timestamp / VI: Timestamp tạo tin nhắn */
  createdAt: string;
  /** EN: Message update timestamp / VI: Timestamp cập nhật tin nhắn */
  updatedAt?: string;
  /** EN: Optional metadata / VI: Metadata tùy chọn */
  metadata?: Record<string, unknown>;
}

/**
 * EN: Conversation interface
 * VI: Interface cuộc trò chuyện
 */
export interface Conversation {
  /** EN: Unique conversation identifier / VI: Mã định danh cuộc trò chuyện duy nhất */
  id: string;
  /** EN: Conversation title / VI: Tiêu đề cuộc trò chuyện */
  title: string;
  /** EN: Last message in conversation / VI: Tin nhắn cuối trong cuộc trò chuyện */
  lastMessage?: ChatMessage;
  /** EN: Last activity timestamp / VI: Timestamp hoạt động cuối */
  updatedAt: string;
  /** EN: Conversation creation timestamp / VI: Timestamp tạo cuộc trò chuyện */
  createdAt: string;
  /** EN: Number of unread messages / VI: Số tin nhắn chưa đọc */
  unreadCount?: number;
}

/**
 * EN: Chat state interface for Zustand store
 * VI: Interface trạng thái chat cho Zustand store
 */
interface ChatState {
  /** EN: WebSocket client instance / VI: Instance WebSocket client */
  wsClient: WebSocketClient | null;
  /** EN: WebSocket connection state / VI: Trạng thái kết nối WebSocket */
  connectionState: WebSocketState;
  /** EN: Current conversation ID / VI: ID cuộc trò chuyện hiện tại */
  currentConversationId: string | null;
  /** EN: List of conversations / VI: Danh sách cuộc trò chuyện */
  conversations: Conversation[];
  /** EN: Messages map by conversation ID / VI: Map tin nhắn theo ID cuộc trò chuyện */
  messages: Record<string, ChatMessage[]>;
  /** EN: Loading state / VI: Trạng thái loading */
  isLoading: boolean;
  /** EN: Error message / VI: Thông báo lỗi */
  error: string | null;
  /** EN: Typing indicator state (userId -> boolean) / VI: Trạng thái chỉ báo đang gõ (userId -> boolean) */
  typingUsers: Record<string, boolean>;
  /** EN: Initialize WebSocket connection / VI: Khởi tạo kết nối WebSocket */
  initializeWebSocket: (url: string) => void;
  /** EN: Connect WebSocket / VI: Kết nối WebSocket */
  connect: () => void;
  /** EN: Disconnect WebSocket / VI: Ngắt kết nối WebSocket */
  disconnect: () => void;
  /** EN: Send message / VI: Gửi tin nhắn */
  sendMessage: (conversationId: string, content: string) => Promise<void>;
  /** EN: Receive message / VI: Nhận tin nhắn */
  receiveMessage: (message: ChatMessage) => void;
  /** EN: Create new conversation / VI: Tạo cuộc trò chuyện mới */
  createConversation: (title?: string) => string;
  /** EN: Select conversation / VI: Chọn cuộc trò chuyện */
  selectConversation: (conversationId: string) => void;
  /** EN: Delete conversation / VI: Xóa cuộc trò chuyện */
  deleteConversation: (conversationId: string) => void;
  /** EN: Update message status / VI: Cập nhật trạng thái tin nhắn */
  updateMessageStatus: (messageId: string, conversationId: string, status: MessageStatus) => void;
  /** EN: Set typing indicator / VI: Đặt chỉ báo đang gõ */
  setTyping: (userId: string, isTyping: boolean) => void;
  /** EN: Clear error / VI: Xóa lỗi */
  clearError: () => void;
}

/**
 * EN: Generate unique ID for messages/conversations
 * VI: Tạo ID duy nhất cho messages/conversations
 */
const generateId = (): string => {
  return `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
};

/**
 * EN: Zustand store for chat state management with WebSocket integration
 * VI: Zustand store để quản lý trạng thái chat với tích hợp WebSocket
 *
 * Features:
 * - WebSocket connection management
 * - Message state management
 * - Conversation state management
 * - Real-time message sending/receiving
 * - Typing indicators
 * - Persistent storage (conversations and messages)
 */
export const useChatStore = create<ChatState>()(
  persist(
    (set, get) => ({
      wsClient: null,
      connectionState: WebSocketState.CLOSED,
      currentConversationId: null,
      conversations: [],
      messages: {},
      isLoading: false,
      error: null,
      typingUsers: {},

      /**
       * EN: Initialize WebSocket client with URL
       * VI: Khởi tạo WebSocket client với URL
       *
       * @param url - WebSocket server URL / URL server WebSocket
       */
      initializeWebSocket: (url: string) => {
        const { wsClient } = get();
        
        // EN: Clean up existing connection if any
        // VI: Dọn dẹp kết nối hiện có nếu có
        if (wsClient) {
          wsClient.disconnect();
        }

        // EN: Create new WebSocket client
        // VI: Tạo WebSocket client mới
        const client = new WebSocketClient({
          url,
          autoReconnect: true,
          maxReconnectAttempts: 10,
          reconnectDelay: 1000,
          callbacks: {
            onOpen: () => {
              set({ connectionState: WebSocketState.CONNECTED });
            },
            onClose: () => {
              set({ connectionState: WebSocketState.CLOSED });
            },
            onError: (error) => {
              console.error('WebSocket error / Lỗi WebSocket:', error);
              set({ 
                connectionState: WebSocketState.ERROR,
                error: 'Connection error / Lỗi kết nối',
              });
            },
            onStateChange: (state) => {
              set({ connectionState: state });
            },
            onMessage: (wsMessage: WebSocketMessage) => {
              const state = get();
              
              // EN: Handle different message types
              // VI: Xử lý các loại tin nhắn khác nhau
              switch (wsMessage.type) {
                case WebSocketMessageType.MESSAGE:
                  if (wsMessage.data) {
                    const chatMessage: ChatMessage = {
                      id: wsMessage.messageId || generateId(),
                      conversationId: wsMessage.conversationId || state.currentConversationId || '',
                      content: (wsMessage.data && typeof wsMessage.data === 'object' && 'content' in wsMessage.data && typeof wsMessage.data.content === 'string') ? wsMessage.data.content : '',
                      sender: (wsMessage.data && typeof wsMessage.data === 'object' && 'sender' in wsMessage.data) ? (wsMessage.data.sender as MessageSender) : MessageSender.ASSISTANT,
                      status: MessageStatus.DELIVERED,
                      createdAt: wsMessage.timestamp || new Date().toISOString(),
                      metadata: (wsMessage.data && typeof wsMessage.data === 'object' && 'metadata' in wsMessage.data) ? (wsMessage.data.metadata as Record<string, unknown>) : undefined,
                    };
                    state.receiveMessage(chatMessage);
                  }
                  break;
                case WebSocketMessageType.TYPING:
                  if (wsMessage.data && typeof wsMessage.data === 'object' && 'userId' in wsMessage.data && typeof wsMessage.data.userId === 'string') {
                    const isTyping = (wsMessage.data && typeof wsMessage.data === 'object' && 'isTyping' in wsMessage.data) ? Boolean(wsMessage.data.isTyping) : false;
                    state.setTyping(wsMessage.data.userId, isTyping);
                  }
                  break;
                case WebSocketMessageType.READ:
                  if (wsMessage.messageId && wsMessage.conversationId) {
                    state.updateMessageStatus(
                      wsMessage.messageId,
                      wsMessage.conversationId,
                      MessageStatus.READ
                    );
                  }
                  break;
                case WebSocketMessageType.ERROR: {
                  const errorMessage = (wsMessage.data && typeof wsMessage.data === 'object' && 'message' in wsMessage.data && typeof wsMessage.data.message === 'string') 
                    ? wsMessage.data.message 
                    : 'Unknown error / Lỗi không xác định';
                  set({ error: errorMessage });
                  break;
                }
                default:
                  console.log('Unhandled WebSocket message type / Loại tin nhắn WebSocket chưa xử lý:', wsMessage.type);
              }
            },
          },
        });

        set({ wsClient: client });
      },

      /**
       * EN: Connect to WebSocket server
       * VI: Kết nối tới server WebSocket
       */
      connect: () => {
        const { wsClient } = get();
        if (!wsClient) {
          throw new Error('WebSocket client not initialized. Call initializeWebSocket first. / WebSocket client chưa được khởi tạo. Gọi initializeWebSocket trước.');
        }
        wsClient.connect();
      },

      /**
       * EN: Disconnect from WebSocket server
       * VI: Ngắt kết nối khỏi server WebSocket
       */
      disconnect: () => {
        const { wsClient } = get();
        if (wsClient) {
          wsClient.disconnect();
        }
        set({ connectionState: WebSocketState.CLOSED });
      },

      /**
       * EN: Send message through WebSocket
       * VI: Gửi tin nhắn qua WebSocket
       *
       * @param conversationId - Conversation ID / ID cuộc trò chuyện
       * @param content - Message content / Nội dung tin nhắn
       */
      sendMessage: async (conversationId: string, content: string) => {
        const { wsClient, currentConversationId } = get();

        // EN: Ensure conversation is selected
        // VI: Đảm bảo cuộc trò chuyện được chọn
        if (currentConversationId !== conversationId) {
          get().selectConversation(conversationId);
        }

        if (!wsClient || !wsClient.isConnected()) {
          throw new Error('WebSocket is not connected / WebSocket chưa kết nối');
        }

        // EN: Create temporary message with SENDING status
        // VI: Tạo tin nhắn tạm thời với trạng thái SENDING
        const tempMessageId = generateId();
        const tempMessage: ChatMessage = {
          id: tempMessageId,
          conversationId,
          content,
          sender: MessageSender.USER,
          status: MessageStatus.SENDING,
          createdAt: new Date().toISOString(),
        };

        // EN: Add temporary message to state
        // VI: Thêm tin nhắn tạm thời vào state
        set((state) => {
          const messages = state.messages[conversationId] || [];
          return {
            messages: {
              ...state.messages,
              [conversationId]: [...messages, tempMessage],
            },
          };
        });

        try {
          // EN: Send message via WebSocket
          // VI: Gửi tin nhắn qua WebSocket
          wsClient.send({
            type: WebSocketMessageType.MESSAGE,
            conversationId,
            messageId: tempMessageId,
            data: {
              content,
              sender: MessageSender.USER,
            },
            timestamp: new Date().toISOString(),
          });

          // EN: Update message status to SENT
          // VI: Cập nhật trạng thái tin nhắn thành SENT
          get().updateMessageStatus(tempMessageId, conversationId, MessageStatus.SENT);
        } catch (error) {
          console.error('Failed to send message / Không thể gửi tin nhắn:', error);
          // EN: Update message status to FAILED
          // VI: Cập nhật trạng thái tin nhắn thành FAILED
          get().updateMessageStatus(tempMessageId, conversationId, MessageStatus.FAILED);
          throw error;
        }
      },

      /**
       * EN: Receive message and add to state
       * VI: Nhận tin nhắn và thêm vào state
       *
       * @param message - Received message / Tin nhắn đã nhận
       */
      receiveMessage: (message: ChatMessage) => {
        set((state) => {
          const messages = state.messages[message.conversationId] || [];
          
          // EN: Check if message already exists (avoid duplicates)
          // VI: Kiểm tra nếu tin nhắn đã tồn tại (tránh trùng lặp)
          const existingIndex = messages.findIndex((m) => m.id === message.id);
          if (existingIndex >= 0) {
            // EN: Update existing message
            // VI: Cập nhật tin nhắn hiện có
            const updatedMessages = [...messages];
            updatedMessages[existingIndex] = message;
            return {
              messages: {
                ...state.messages,
                [message.conversationId]: updatedMessages,
              },
            };
          }

          // EN: Add new message
          // VI: Thêm tin nhắn mới
          return {
            messages: {
              ...state.messages,
              [message.conversationId]: [...messages, message],
            },
          };
        });

        // EN: Update conversation's last message
        // VI: Cập nhật tin nhắn cuối của cuộc trò chuyện
        set((state) => {
          const conversations = state.conversations.map((conv) => {
            if (conv.id === message.conversationId) {
              return {
                ...conv,
                lastMessage: message,
                updatedAt: message.createdAt,
                unreadCount: conv.id === state.currentConversationId 
                  ? 0 
                  : (conv.unreadCount || 0) + 1,
              };
            }
            return conv;
          });

          // EN: Create conversation if it doesn't exist
          // VI: Tạo cuộc trò chuyện nếu chưa tồn tại
          if (!conversations.find((c) => c.id === message.conversationId)) {
            const newConversation: Conversation = {
              id: message.conversationId,
              title: message.content.substring(0, 50) || 'New Conversation / Cuộc trò chuyện mới',
              lastMessage: message,
              updatedAt: message.createdAt,
              createdAt: message.createdAt,
              unreadCount: message.conversationId === state.currentConversationId ? 0 : 1,
            };
            conversations.unshift(newConversation);
          }

          return { conversations };
        });
      },

      /**
       * EN: Create new conversation
       * VI: Tạo cuộc trò chuyện mới
       *
       * @param title - Optional conversation title / Tiêu đề cuộc trò chuyện tùy chọn
       * @returns New conversation ID / ID cuộc trò chuyện mới
       */
      createConversation: (title?: string) => {
        const conversationId = generateId();
        const now = new Date().toISOString();
        
        const newConversation: Conversation = {
          id: conversationId,
          title: title || 'New Conversation / Cuộc trò chuyện mới',
          updatedAt: now,
          createdAt: now,
          unreadCount: 0,
        };

        set((state) => ({
          conversations: [newConversation, ...state.conversations],
          currentConversationId: conversationId,
          messages: {
            ...state.messages,
            [conversationId]: [],
          },
        }));

        return conversationId;
      },

      /**
       * EN: Select conversation and mark messages as read
       * VI: Chọn cuộc trò chuyện và đánh dấu tin nhắn đã đọc
       *
       * @param conversationId - Conversation ID to select / ID cuộc trò chuyện cần chọn
       */
      selectConversation: (conversationId: string) => {
        set((state) => {
          // EN: Update conversations to clear unread count
          // VI: Cập nhật cuộc trò chuyện để xóa số tin nhắn chưa đọc
          const conversations = state.conversations.map((conv) => {
            if (conv.id === conversationId) {
              return { ...conv, unreadCount: 0 };
            }
            return conv;
          });

          return {
            currentConversationId: conversationId,
            conversations,
          };
        });
      },

      /**
       * EN: Delete conversation and its messages
       * VI: Xóa cuộc trò chuyện và tin nhắn của nó
       *
       * @param conversationId - Conversation ID to delete / ID cuộc trò chuyện cần xóa
       */
      deleteConversation: (conversationId: string) => {
        set((state) => {
          const conversations = state.conversations.filter((c) => c.id !== conversationId);
          const messages = { ...state.messages };
          delete messages[conversationId];

          // EN: Clear current conversation if it was deleted
          // VI: Xóa cuộc trò chuyện hiện tại nếu nó bị xóa
          const currentConversationId = state.currentConversationId === conversationId 
            ? null 
            : state.currentConversationId;

          return {
            conversations,
            messages,
            currentConversationId,
          };
        });
      },

      /**
       * EN: Update message status
       * VI: Cập nhật trạng thái tin nhắn
       *
       * @param messageId - Message ID / ID tin nhắn
       * @param conversationId - Conversation ID / ID cuộc trò chuyện
       * @param status - New status / Trạng thái mới
       */
      updateMessageStatus: (messageId: string, conversationId: string, status: MessageStatus) => {
        set((state) => {
          const messages = state.messages[conversationId] || [];
          const updatedMessages = messages.map((msg) => {
            if (msg.id === messageId) {
              return {
                ...msg,
                status,
                updatedAt: new Date().toISOString(),
              };
            }
            return msg;
          });

          return {
            messages: {
              ...state.messages,
              [conversationId]: updatedMessages,
            },
          };
        });
      },

      /**
       * EN: Set typing indicator for user
       * VI: Đặt chỉ báo đang gõ cho người dùng
       *
       * @param userId - User ID / ID người dùng
       * @param isTyping - Typing state / Trạng thái đang gõ
       */
      setTyping: (userId: string, isTyping: boolean) => {
        set((state) => ({
          typingUsers: {
            ...state.typingUsers,
            [userId]: isTyping,
          },
        }));
      },

      /**
       * EN: Clear error message
       * VI: Xóa thông báo lỗi
       */
      clearError: () => {
        set({ error: null });
      },
    }),
    {
      // EN: Persist chat state to localStorage
      // VI: Persist trạng thái chat vào localStorage
      name: 'chat-storage',
      // EN: Only persist conversations and messages, exclude WebSocket client and connection state
      // VI: Chỉ persist conversations và messages, loại trừ WebSocket client và connection state
      partialize: (state) => ({
        conversations: state.conversations,
        messages: state.messages,
        currentConversationId: state.currentConversationId,
      }),
    }
  )
);
