import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
import {
  type AiStreamChunkPayload,
  type RoomMessagePayload,
  type RoomWsStatus,
  type WsOutEvent,
} from '@/lib/room';
import { requestWsToken, buildWsUrlWithToken } from '@/lib/ws-token';
import { client } from '@/client/client.gen';
import type { AxiosResponse } from 'axios';
import type { RoomMemberResponse } from '@/client';

const RECONNECT_BASE_DELAY = 1_000;
const RECONNECT_MAX_DELAY = 15_000;

/** A message as held in the UI state */
export type UiMessage = RoomMessagePayload & {
  /** Set while the server is streaming an AI reply into this message */
  is_streaming?: boolean;
  /** Accumulated streaming content; flushed to content on `done: true` */
  streaming_content?: string;
  /** Display name resolved from sender_id; undefined if not yet resolved */
  display_name?: string;
  /** Avatar URL resolved from members list */
  avatar_url?: string;
  /** For optimistic UI: message failed to send */
  isOptimisticError?: boolean;
  /** Reply to message ID */
  in_reply_to?: string | null;
  /** Edited timestamp */
  edited_at?: string | null;
  /** Revoked timestamp */
  revoked?: string | null;
  /** Revoked by user ID */
  revoked_by?: string | null;
};

type RoomMessageCacheEntry = {
  messages: UiMessage[];
  isHistoryLoaded: boolean;
  /** seq of the latest message, used as cursor for pagination */
  nextCursor: number | null;
};

interface MessageListResponse {
  code: number;
  message: string;
  data: { messages: RestMessage[]; total: number };
}

/** REST message shape (matches RoomMessageResponse from the SDK) */
interface RestMessage {
  id: string;
  seq: number;
  room: string;
  sender_type: string;
  sender_id?: string | null;
  display_name?: string | null;
  thread?: string | null;
  in_reply_to?: string | null;
  content: string;
  content_type: string;
  edited_at?: string | null;
  send_at: string;
  revoked?: string | null;
  revoked_by?: string | null;
}

/** Display name and avatar URL resolved from a message's sender */
interface SenderInfo {
  displayName: string;
  avatarUrl: string | undefined;
}

/** Resolve displayName and avatar URL for a message sender.
 * - AI messages: use sender_id (the model UUID)
 * - User messages: look up user_info from members list, fall back to sender_id */
function resolveSender(payload: RoomMessagePayload, members: RoomMemberResponse[]): SenderInfo {
  if (payload.sender_type === 'ai') {
    return { displayName: payload.sender_id ?? 'AI', avatarUrl: undefined };
  }
  if (payload.sender_id) {
    const member = members.find((m) => m.user === payload.sender_id);
    if (member) {
      const username = member.user_info?.username ?? member.user;
      return { displayName: username, avatarUrl: member.user_info?.avatar_url ?? undefined };
    }
  }
  if (payload.sender_type === 'system') return { displayName: 'System', avatarUrl: undefined };
  return { displayName: payload.sender_type, avatarUrl: undefined };
}

function compareMessages(a: UiMessage, b: UiMessage): number {
  const timeDiff = new Date(a.send_at).getTime() - new Date(b.send_at).getTime();
  return timeDiff !== 0 ? timeDiff : a.id.localeCompare(b.id);
}

function insertSorted(arr: UiMessage[], msg: UiMessage): UiMessage[] {
  const result = [...arr];
  let lo = 0;
  let hi = result.length;
  while (lo < hi) {
    const mid = (lo + hi) >>> 1;
    if (compareMessages(result[mid], msg) < 0) lo = mid + 1;
    else hi = mid;
  }
  result.splice(lo, 0, msg);
  return result;
}

export interface UseRoomWsOptions {
  /** VITE_API_BASE_URL value (without /ws suffix) */
  baseUrl: string;
  /** Currently open room ID */
  roomId: string | null;
  /** Limit for initial history load */
  historyLimit?: number;
  /** Room members, used to resolve display_name for user messages */
  members?: RoomMemberResponse[];
  /** Called when the AI streaming chunk arrives */
  onAiStreamChunk?: (payload: AiStreamChunkPayload) => void;
}

export interface UseRoomWsReturn {
  messages: UiMessage[];
  status: RoomWsStatus;
  errorMessage: string | null;
  isHistoryLoaded: boolean;
  isLoadingMore: boolean;
  nextCursor: number | null;
  /** Load older messages (called when user scrolls to top) */
  loadMore: (cursor?: number | null) => void;
}

/**
 * Manages a WebSocket connection for a single room.
 *
 * Features:
 * - Auto-reconnect with exponential back-off
 * - Per-room message cache so switching rooms preserves scroll position
 * - AI streaming chunk accumulation via `streaming_content`
 * - `loadMore` for cursor-based history pagination
 */
export function useRoomWs({
  baseUrl,
  roomId,
  historyLimit = 50,
  members = [],
  onAiStreamChunk,
}: UseRoomWsOptions): UseRoomWsReturn {
  const [messages, setMessages] = useState<UiMessage[]>([]);
  const [status, setStatus] = useState<RoomWsStatus>('idle');
  const [errorMessage, setErrorMessage] = useState<string | null>(null);
  const [isHistoryLoaded, setIsHistoryLoaded] = useState(false);
  const [isLoadingMore, setIsLoadingMore] = useState(false);
  const [nextCursor, setNextCursor] = useState<number | null>(null);

  const wsRef = useRef<WebSocket | null>(null);
  const activeRoomIdRef = useRef<string | null>(null);
  const shouldReconnectRef = useRef(true);
  const reconnectAttemptRef = useRef(0);
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const wsTokenRef = useRef<string | null>(null);

  const roomCacheRef = useRef<Map<string, RoomMessageCacheEntry>>(new Map());
  /** Ref to current messages for use inside event handlers */
  const messagesRef = useRef<UiMessage[]>([]);
  messagesRef.current = messages;

  /** Ref to current nextCursor */
  const nextCursorRef = useRef<number | null>(null);
  nextCursorRef.current = nextCursor;

  /** Ref to members, used for display_name resolution */
  const membersRef = useRef<RoomMemberResponse[]>([]);
  membersRef.current = members;

  /** Ref for AI streaming RAF batch */
  const streamingBatchRef = useRef<Map<string, { content: string; done: boolean }>>(new Map());
  const streamingRafRef = useRef<number | null>(null);

  /** Flush streaming batch to state */
  const flushStreamingBatch = useCallback(() => {
    const batch = streamingBatchRef.current;
    if (batch.size === 0) return;

    setMessages((prev) => {
      const next = [...prev];
      let changed = false;

      for (const [messageId, chunk] of batch) {
        const idx = next.findIndex((m) => m.id === messageId);
        if (idx === -1) {
          const placeholder: UiMessage = {
            id: messageId,
            room_id: chunk.content.split('\n')[0]?.slice(0, 100) || '',
            sender_type: 'ai',
            content: chunk.done ? chunk.content : '',
            content_type: 'text',
            send_at: new Date().toISOString(),
            seq: 0,
            display_name: 'AI',
            is_streaming: !chunk.done,
            streaming_content: chunk.done ? undefined : chunk.content,
          };
          next.push(placeholder);
          changed = true;
        } else {
          const updated = { ...next[idx] };
          if (chunk.done) {
            updated.is_streaming = false;
            updated.content = chunk.content;
            updated.streaming_content = undefined;
          } else {
            updated.is_streaming = true;
            updated.streaming_content = (updated.streaming_content ?? updated.content) + chunk.content;
          }
          next[idx] = updated;
          changed = true;
        }
      }

      return changed ? next : prev;
    });

    streamingBatchRef.current.clear();
    streamingRafRef.current = null;
  }, []);

  useEffect(() => {
    const room = activeRoomIdRef.current;
    if (!room) return;
    roomCacheRef.current.set(room, {
      messages,
      isHistoryLoaded,
      nextCursor,
    });
  }, [messages, isHistoryLoaded, nextCursor]);

  const connectWs = useCallback(
    async (roomUid: string) => {
      if (!shouldReconnectRef.current || activeRoomIdRef.current !== roomUid) return;

      // Build URL with token if available
      const url = buildWsUrlWithToken(baseUrl, `/ws/rooms/${roomUid}`, wsTokenRef.current);
      console.debug('[useRoomWs] connecting to', url, { baseUrl, roomUid });

      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
        reconnectTimeoutRef.current = null;
      }

      setStatus('connecting');
      setErrorMessage(null);

      const ws = new WebSocket(url);
      wsRef.current = ws;

      ws.onopen = () => {
        if (activeRoomIdRef.current !== roomUid) return;
        reconnectAttemptRef.current = 0;
        setStatus('open');
        setErrorMessage(null);
        console.debug('[useRoomWs] ws opened for room', roomUid);
      };

      ws.onmessage = (ev: MessageEvent<string>) => {
        if (activeRoomIdRef.current !== roomUid) return;

        let event: WsOutEvent;
        try {
          event = JSON.parse(ev.data) as WsOutEvent;
        } catch {
          console.warn('[useRoomWs] parse error, data:', ev.data);
          setErrorMessage('Invalid WebSocket message');
          return;
        }

        if ('error' in event && event.error) {
          console.warn('[useRoomWs] error event:', event.error);
          setErrorMessage(event.error);
          return;
        }

        if (!('event' in event) || !event.event) {
          console.warn('[useRoomWs] no event field, raw:', event);
          return;
        }

        console.debug('[useRoomWs] received event type:', event.event.type, event.event);

        switch (event.event.type) {
          case 'room_message': {
            // Backend sends payload flat on event.event (no data wrapper); also support nested data
            const raw = (event.event as any);
            const incoming: RoomMessagePayload = raw.data ?? raw;
            console.debug('[useRoomWs] room_message:', incoming.id, incoming.content);

            // Use Set for O(1) duplicate check instead of O(n) Array.some
            const existingIds = new Set(messagesRef.current.map((m) => m.id));
            if (existingIds.has(incoming.id)) {
              console.debug('[useRoomWs] duplicate message, skipping');
              return;
            }

            const sender = resolveSender(incoming, membersRef.current);
            const display_name = incoming.display_name ?? sender.displayName;
            const avatar_url = sender.avatarUrl;

            setMessages((prev) =>
              insertSorted(prev, { ...incoming, display_name, avatar_url }),
            );
            break;
          }

          case 'ai_stream_chunk': {
            const raw = event.event as any;
            const chunk: AiStreamChunkPayload = raw.data ?? raw;
            onAiStreamChunk?.(chunk);

            // Batch streaming chunks using RAF to reduce re-render frequency
            streamingBatchRef.current.set(chunk.message_id, {
              content: chunk.content,
              done: chunk.done,
            });

            if (streamingRafRef.current == null) {
              streamingRafRef.current = requestAnimationFrame(() => {
                flushStreamingBatch();
              });
            }
            break;
          }

          default:
            break;
        }
      };

      ws.onclose = (ev: CloseEvent) => {
        console.debug('[useRoomWs] WebSocket closed', { code: ev.code, reason: ev.reason, wasClean: ev.wasClean });
        const activeSocket = wsRef.current;
        if (activeSocket !== ws) return;
        wsRef.current = null;
        if (activeRoomIdRef.current !== roomUid) return;

        setStatus('closed');
        if (shouldReconnectRef.current) {
          const attempt = ++reconnectAttemptRef.current;
          const delay = Math.min(RECONNECT_BASE_DELAY * 2 ** (attempt - 1), RECONNECT_MAX_DELAY);
          reconnectTimeoutRef.current = setTimeout(() => connectWs(roomUid), delay);
        }
      };

      ws.onerror = (ev: Event) => {
        console.error('[useRoomWs] WebSocket error', ev);
        if (activeRoomIdRef.current !== roomUid) return;
        setErrorMessage('WebSocket error');
      };
    },
    [baseUrl, onAiStreamChunk],
  );

  useEffect(() => {
    const prevRoom = activeRoomIdRef.current;

    if (!roomId) {
      // Disconnect
      activeRoomIdRef.current = null;
      shouldReconnectRef.current = false;
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
        reconnectTimeoutRef.current = null;
      }
      if (wsRef.current) {
        wsRef.current.close();
        wsRef.current = null;
      }
      setMessages([]);
      setStatus('idle');
      setErrorMessage(null);
      setIsHistoryLoaded(false);
      setNextCursor(null);
      return;
    }

    // Save previous room's state
    if (prevRoom && prevRoom !== roomId) {
      roomCacheRef.current.set(prevRoom, {
        messages: messagesRef.current,
        isHistoryLoaded,
        nextCursor: nextCursorRef.current,
      });
    }

    activeRoomIdRef.current = roomId;
    shouldReconnectRef.current = true;
    reconnectAttemptRef.current = 0;

    // Fetch WS token before connecting
    const connectWithToken = async () => {
      try {
        const token = await requestWsToken();
        wsTokenRef.current = token;
      } catch (error) {
        console.warn('[useRoomWs] Failed to fetch WS token, falling back to cookie auth:', error);
        wsTokenRef.current = null;
      }

      // Restore from cache or start fresh
      const cached = roomCacheRef.current.get(roomId);
      if (cached) {
        setMessages(cached.messages);
        setIsHistoryLoaded(cached.isHistoryLoaded);
        setNextCursor(cached.nextCursor);
      } else {
        setMessages([]);
        setIsHistoryLoaded(false);
        setNextCursor(null);
        // Load initial history via REST (WS is push-only, can't request history)
        if (roomId) {
          client
            .get({ url: `/api/rooms/${roomId}/messages`, params: { limit: historyLimit } })
            .then((resp) => {
              const r = resp as AxiosResponse<MessageListResponse>;
              if (activeRoomIdRef.current !== roomId) return;
              const msgs = (r.data?.data?.messages ?? []).map((m) => {
                const sender = resolveSender({ ...m, room_id: m.room } as RoomMessagePayload, members);
                const display_name = m.display_name ?? sender.displayName;
                const avatar_url = sender.avatarUrl;
                return {
                  ...m,
                  room_id: m.room,
                  thread_id: m.thread ?? null,
                  display_name,
                  avatar_url,
                };
              });
              setMessages(msgs);
              setNextCursor(msgs.length > 0 ? msgs[msgs.length - 1].seq : null);
              setIsHistoryLoaded(msgs.length < historyLimit);
            })
            .catch(() => {
              if (activeRoomIdRef.current !== roomId) return;
              setIsHistoryLoaded(true);
            });
        }
      }

      // Close other connections (shouldn't be any in practice)
      if (wsRef.current) {
        wsRef.current.close();
        wsRef.current = null;
      }

      connectWs(roomId);
    };

    connectWithToken();

    return () => {
      // Save state before unmounting
      const room = activeRoomIdRef.current;
      if (room) {
        roomCacheRef.current.set(room, {
          messages: messagesRef.current,
          isHistoryLoaded: isHistoryLoaded,
          nextCursor: nextCursorRef.current,
        });
      }
      shouldReconnectRef.current = false;
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
        reconnectTimeoutRef.current = null;
      }
      if (wsRef.current) {
        wsRef.current.close();
        wsRef.current = null;
      }
      if (streamingRafRef.current != null) {
        cancelAnimationFrame(streamingRafRef.current);
        streamingRafRef.current = null;
      }
      activeRoomIdRef.current = null;
    };
  }, [roomId, connectWs]);

  const loadMore = useCallback(
    async (cursor?: number | null) => {
      if (!roomId || isLoadingMore) return;
      // Use REST API for history pagination — WS is push-only
      const effectiveCursor = cursor ?? nextCursorRef.current;
      if (effectiveCursor == null) return;

      setIsLoadingMore(true);
      try {
        const resp = await client.get({ url: `/api/rooms/${roomId}/messages`, params: { before_seq: effectiveCursor, limit: historyLimit } }) as AxiosResponse<MessageListResponse>;
        const older = (resp.data?.data?.messages ?? []).map((m) => {
          const sender = resolveSender({ ...m, room_id: m.room } as RoomMessagePayload, membersRef.current);
          const display_name = m.display_name ?? sender.displayName;
          const avatar_url = sender.avatarUrl;
          return {
            ...m,
            room_id: m.room,
            thread_id: m.thread ?? null,
            display_name,
            avatar_url,
          };
        });
        if (older.length === 0) {
          setIsHistoryLoaded(true);
          return;
        }
        // Prepend older messages (they arrive in ascending seq order)
        setMessages((prev) => {
          const existingIds = new Set(prev.map((m) => m.id));
          const newOnes = older.filter((m) => !existingIds.has(m.id));
          if (newOnes.length === 0) {
            setIsHistoryLoaded(true);
            return prev;
          }
          // New cursor = smallest seq among loaded messages
          const newCursor = newOnes[newOnes.length - 1].seq;
          setNextCursor(newCursor > 0 ? newCursor : null);
          return [...newOnes, ...prev];
        });
      } catch {
        // Non-critical — just stop trying
        setIsHistoryLoaded(true);
      } finally {
        setIsLoadingMore(false);
      }
    },
    [roomId, historyLimit, isLoadingMore],
  );

  return useMemo(
    () => ({
      messages,
      status,
      errorMessage,
      isHistoryLoaded,
      isLoadingMore,
      nextCursor,
      loadMore,
    }),
    [messages, status, errorMessage, isHistoryLoaded, isLoadingMore, nextCursor, loadMore],
  );
}
