import { useEffect, useRef, useCallback, useState } from 'react';
import { useChatStore } from '@/store/chatStore';
import { useSessionStore } from '@/store/sessionStore';
import { useNotificationStore } from '@/store/notificationStore';
import { getToken } from '@/api/client';
import type { Message, WSClientMessage, WSServerMessage, MessageDisplayType } from '@/types/chat';

function getWsUrl(): string {
  if (import.meta.env.VITE_WS_URL) return import.meta.env.VITE_WS_URL;
  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
  return `${protocol}//${window.location.host}/ws/chat`;
}

// ========== 指数退避重连配置 ==========
const INITIAL_RECONNECT_DELAY_MS = 1000; // 初始延迟 1 秒
const MAX_RECONNECT_DELAY_MS = 30000;     // 最大延迟 30 秒
const MAX_RECONNECT_ATTEMPTS = 10;        // 最大重连次数
const RECONNECT_MULTIPLIER = 2;           // 每次翻倍

/**
 * 计算带 jitter 的指数退避延迟
 * delay = min(initial * multiplier^attempt, maxDelay)
 * jitter = delay * random(0.5, 1.0)，避免惊群效应
 */
function getBackoffDelay(attempt: number): number {
  const exponentialDelay = INITIAL_RECONNECT_DELAY_MS * Math.pow(RECONNECT_MULTIPLIER, attempt);
  const cappedDelay = Math.min(exponentialDelay, MAX_RECONNECT_DELAY_MS);
  // 添加 jitter：在 [delay/2, delay] 范围内随机
  const jitter = cappedDelay * (0.5 + Math.random() * 0.5);
  return Math.floor(jitter);
}

let wsInstanceCounter = 0;

// ========== 跨端消息去重 ==========
const sentMsgIds = new Set<string>();
const MAX_SENT_IDS = 40;

function trackSentMsgId(id: string) {
  sentMsgIds.add(id);
  if (sentMsgIds.size > MAX_SENT_IDS) {
    const entries = [...sentMsgIds];
    sentMsgIds.clear();
    for (const entry of entries.slice(-MAX_SENT_IDS / 2)) {
      sentMsgIds.add(entry);
    }
  }
}

// 会话作用域的消息类型：仅处理匹配当前会话的消息
const SESSION_SCOPED_TYPES = new Set([
  'response', 'stream_start', 'stream_chunk', 'stream_end',
  'stream_segments', 'review', 'multi_message', 'thinking',
  'tool_progress', 'system_info',
]);

export function useWebSocket() {
  const [isConnected, setIsConnected] = useState(false);
  const [reconnectAttempts, setReconnectAttempts] = useState(0);
  const [maxReconnectAttempts] = useState(MAX_RECONNECT_ATTEMPTS);

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const shouldReconnectRef = useRef(true);
  const activeSessionRef = useRef<string | null>(null);
  const reconnectCountRef = useRef(0);
  const instanceIdRef = useRef(++wsInstanceCounter);

  // 订阅 sessionStore 中的 currentSessionId 变化
  const currentSessionId = useSessionStore((s) => s.currentSessionId);

  const connect = useCallback(() => {
    const instanceId = instanceIdRef.current;
    const token = getToken();
    if (!token) {
      console.warn(`[WS#${instanceId}] connect: 无 token，跳过连接`);
      return;
    }

    // 检查是否超过最大重连次数
    if (reconnectCountRef.current >= MAX_RECONNECT_ATTEMPTS) {
      console.error(`[WS#${instanceId}] 已达到最大重连次数 (${MAX_RECONNECT_ATTEMPTS})，停止重连`);
      setReconnectAttempts(reconnectCountRef.current);
      // 通知用户连接失败
      useChatStore.getState().setTyping(false);
      useChatStore.getState().addMessage({
        id: 'err_max_reconnect_' + Date.now(),
        role: 'system',
        content: `⚠️ WebSocket 连接失败，已尝试重连 ${MAX_RECONNECT_ATTEMPTS} 次，请刷新页面后重试`,
        timestamp: Date.now(),
      });
      return;
    }

    // 关闭旧连接
    if (wsRef.current) {
      console.log(`[WS#${instanceId}] 关闭旧连接`);
      shouldReconnectRef.current = false;
      wsRef.current.close();
      wsRef.current = null;
    }

    const sessionID = useSessionStore.getState().currentSessionId || '';
    const url = sessionID
      ? `${getWsUrl()}?token=${token}&session_id=${sessionID}`
      : `${getWsUrl()}?token=${token}`;

    console.log(`[WS#${instanceId}] 正在连接, session_id=${sessionID || '(无)'}, reconnectAttempt=${reconnectCountRef.current}`);
    const ws = new WebSocket(url);

    ws.onopen = () => {
      setIsConnected(true);
      shouldReconnectRef.current = true;
      // 连接成功后重置退避计数器
      reconnectCountRef.current = 0;
      setReconnectAttempts(0);
      console.log(`[WS#${instanceId}] 已连接, session_id:`, sessionID);

      // 连接后发送会话恢复消息，恢复后端上下文
      const sid = useSessionStore.getState().currentSessionId;
      if (sid) {
        const resumeMsg: WSClientMessage = {
          type: 'history',
          session_id: sid,
          timestamp: Date.now(),
        };
        ws.send(JSON.stringify(resumeMsg));
        console.log(`[WS#${instanceId}] 已发送会话恢复请求, session_id:`, sid);
      }
    };

    ws.onclose = () => {
      setIsConnected(false);
      console.log(`[WS#${instanceId}] 已断开`);
      if (shouldReconnectRef.current) {
        const currentAttempt = reconnectCountRef.current;
        const delay = getBackoffDelay(currentAttempt);
        console.log(`[WS#${instanceId}] ${delay}ms 后重连 (第 ${currentAttempt + 1}/${MAX_RECONNECT_ATTEMPTS} 次)...`);
        if (reconnectTimerRef.current) {
          clearTimeout(reconnectTimerRef.current);
        }
        reconnectTimerRef.current = setTimeout(() => {
          reconnectCountRef.current += 1;
          setReconnectAttempts(reconnectCountRef.current);
          connect();
        }, delay);
      }
    };

    ws.onerror = (err) => {
      console.error(`[WS#${instanceId}] 连接错误:`, err);
    };

    ws.onmessage = (event) => {
      try {
        const msg: WSServerMessage = JSON.parse(event.data);
        handleServerMessage(msg);
      } catch (err) {
        console.error(`[WS#${instanceId}] 消息解析失败:`, err);
      }
    };

    wsRef.current = ws;
  }, []);

  // 会话切换时：重建 WebSocket 连接（消息历史由 useSession.setCurrentSession 负责加载）
  useEffect(() => {
    activeSessionRef.current = currentSessionId;

    // 重置重连计数（切换会话时是新连接）
    reconnectCountRef.current = 0;
    setReconnectAttempts(0);

    connect();

    return () => {
      if (reconnectTimerRef.current) {
        clearTimeout(reconnectTimerRef.current);
      }
      shouldReconnectRef.current = false;
      wsRef.current?.close();
    };
  }, [connect, currentSessionId]);

  const sendMessage = useCallback((msg: WSClientMessage) => {
    const instanceId = instanceIdRef.current;
    // 注册 client_msg_id 用于跨端回显去重
    if (msg.client_msg_id) {
      trackSentMsgId(msg.client_msg_id);
    }
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      const sessionID = useSessionStore.getState().currentSessionId;
      wsRef.current.send(
        JSON.stringify({
          ...msg,
          session_id: msg.session_id || sessionID || undefined,
        })
      );
    } else {
      const stateLabels = ['CONNECTING', 'OPEN', 'CLOSING', 'CLOSED'];
      const state = wsRef.current
        ? (stateLabels[wsRef.current.readyState] || `UNKNOWN(${wsRef.current.readyState})`)
        : 'NO_SOCKET';
      console.error(
        `[WS#${instanceId}] sendMessage 失败: WebSocket 未就绪 (readyState=${state})，消息被丢弃:`,
        msg.type,
        msg.content?.slice(0, 50)
      );
      // 通知用户连接已断开
      useChatStore.getState().setTyping(false);
      useChatStore.getState().addMessage({
        id: 'err_' + Date.now(),
        role: 'system',
        content: '⚠️ 连接未建立，请刷新页面后重试',
        timestamp: Date.now(),
      });
    }
  }, []);

  return { isConnected, sendMessage, reconnectAttempts, maxReconnectAttempts };
}

function handleServerMessage(msg: WSServerMessage) {
  const { addMessage, appendToLastMessage, finishStreaming, setTyping } =
    useChatStore.getState();
  const { setMessages } = useSessionStore.getState();
  const chatState = useChatStore.getState();

  // 会话作用域消息：仅处理匹配当前会话的消息，避免跨会话串消息
  if (SESSION_SCOPED_TYPES.has(msg.type) && msg.session_id) {
    const currentSid = useSessionStore.getState().currentSessionId;
    if (currentSid && msg.session_id !== currentSid) {
      return;
    }
  }

  switch (msg.type) {
    case 'stream_start':
	      setTyping(true);
	      break;

    case 'response':
      if (msg.text || msg.content) {
        // 跨端用户消息去重：本端已本地添加，跳过服务端回显
        if (msg.role === 'user' && msg.message_id && sentMsgIds.has(msg.message_id)) {
          break;
        }
        addMessage({
          id: msg.message_id || '',
          role: (msg.role as Message['role']) || 'assistant',
          content: (msg.text || msg.content) as string,
          timestamp: msg.timestamp,
          msgType: (msg.msg_type as MessageDisplayType) || undefined,
          client_info: msg.client_info,
          audioUrl: msg.full_audio_url,
          segments: msg.segments,
          metadata: msg.metadata || (msg.tool_calls ? { tool_calls: msg.tool_calls } : undefined),
        });
      }
      setTyping(false);
      break;

    case 'stream_chunk':
      if (msg.content) {
        const { messages } = useChatStore.getState();
        const lastMsg = messages[messages.length - 1];
        if (
          !lastMsg ||
          lastMsg.role !== 'assistant' ||
          !lastMsg.isStreaming
        ) {
          addMessage({
            id: msg.message_id || 'msg_' + Date.now(),
            role: 'assistant',
            content: msg.content,
            timestamp: msg.timestamp,
            isStreaming: true,
            msgType: 'chat',
          });
          setTyping(false);
        } else {
          appendToLastMessage(msg.content);
        }
      }
      break;

    case 'stream_end':
      finishStreaming();
      break;

    case 'history_response':
      // 防御性检查：仅当当前消息为空时才加载 WebSocket 历史响应
      if (msg.messages && msg.messages.length > 0) {
        const sessionState = useSessionStore.getState();
        if (sessionState.messages.length > 0 || chatState.messages.length > 0) {
          console.log(
            '[WS] 忽略 history_response：消息已由 HTTP 加载',
            `sessionMessages=${sessionState.messages.length}`,
            `chatMessages=${chatState.messages.length}`
          );
          break;
        }
        const msgsWithIds: Message[] = msg.messages.map((m, i) => ({
          ...m,
          id: m.id || `hist_${i}_${Date.now()}`,
          msgType: m.msgType || 'chat',
        }));
        setMessages(msgsWithIds);
        useChatStore.getState().setMessages(msgsWithIds);
      } else if (msg.messages && msg.messages.length === 0) {
        console.log('[WS] 忽略空的 history_response（后端缓存和数据库均无消息），等待 HTTP 加载');
      }
      setTyping(false);
      break;

    case 'review':
      // 审查子会话消息 — 后端返回带类型的 review_messages 列表
      // 支持类型: action, chat, markdown, code, search_result
      if (msg.review_messages && msg.review_messages.length > 0) {
        for (const rm of msg.review_messages) {
          const msgType = (rm.type as MessageDisplayType) || 'chat';
          addMessage({
            id: `review_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`,
            role: rm.type === 'action' ? 'action' : 'assistant',
            content: rm.content,
            timestamp: Date.now(),
            isStreaming: false,
            msgType,
            metadata: rm.metadata,
          } as Message);
        }
      }
      setTyping(false);
      break;

    case 'multi_message':
      // 多段消息 — 后端已分类 msg_type，前端直接使用
      if (msg.multi_messages && msg.multi_messages.length > 0) {
        for (const item of msg.multi_messages) {
          const itemMsgType = (item.msg_type as MessageDisplayType) || 'chat';
          addMessage({
            id: `multi_${Date.now()}_${item.index}`,
            role: itemMsgType === 'action' ? 'action' : 'assistant',
            content: item.content,
            timestamp: msg.timestamp || Date.now(),
            isStreaming: false,
            msgType: itemMsgType,
          });
        }
      }
      setTyping(false);
      break;

    case 'stream_segments':
      // 流式片段 — 更新最后一条助手消息的 segments 和 full_audio_url
      if (msg.segments || msg.full_audio_url) {
        const { messages: currentMsgs } = useChatStore.getState();
        if (currentMsgs.length > 0) {
          const lastMsg = currentMsgs[currentMsgs.length - 1];
          if (lastMsg.role === 'assistant') {
            const updated = { ...lastMsg };
            if (msg.segments) updated.segments = msg.segments;
            if (msg.full_audio_url) updated.audioUrl = msg.full_audio_url;
            useChatStore.getState().updateMessage(lastMsg.id, updated);
          }
        }
      }
      break;

    case 'device_update':
      if (msg.devices && msg.devices.length > 0) {
        chatState.setIoTDevices(msg.devices);
      }
      break;

    case 'background_thinking':
      if (msg.thinking_status) {
        chatState.setBackgroundThinkingStatus(msg.thinking_status);
      }
      break;

    case 'notification':
      if (msg.notification) {
        // 添加到通知 store
        const { addNotification } = useNotificationStore.getState();
        addNotification(msg.notification);

        // 触发浏览器桌面通知
        if (typeof Notification !== 'undefined') {
          if (Notification.permission === 'granted') {
            const n = new Notification(msg.notification.title, {
              body: msg.notification.body,
              icon: '/images/Cyrene_Avatar/3rd_Form/Cyrene-3F-Q-Happy-1.png',
              tag: msg.notification.id,
              data: msg.notification.data,
            });
            n.onclick = () => {
              window.focus();
              if (msg.notification?.data?.session_id) {
                const { setCurrentSessionId } = useSessionStore.getState();
                const sid = msg.notification.data.session_id as string;
                if (sid) setCurrentSessionId(sid);
              }
              n.close();
            };
          } else if (Notification.permission === 'default') {
            Notification.requestPermission();
          }
        }
      }
      break;

    case 'error':
      console.error('[WS] 服务端错误:', msg.error);
      setTyping(false);
      // 在聊天界面显示错误提示，让用户知情
      addMessage({
        id: 'err_' + Date.now(),
        role: 'system',
        content: '❌ 服务端错误: ' + (msg.error || '未知错误'),
        timestamp: msg.timestamp || Date.now(),
      });
      break;

    case 'thinking':
      if (msg.thinking_content) {
        addMessage({
          id: msg.message_id || 'think_' + Date.now(),
          role: 'assistant',
          content: msg.thinking_content,
          timestamp: msg.timestamp || Date.now(),
          msgType: 'thinking',
          isStreaming: false,
        });
      }
      break;

    case 'tool_progress':
      if (msg.tool_progress) {
        addMessage({
          id: msg.message_id || 'tool_' + Date.now(),
          role: 'system',
          content: `${msg.tool_progress.tool_name}: ${msg.tool_progress.message}`,
          timestamp: msg.timestamp || Date.now(),
          msgType: 'tool_progress',
          isStreaming: false,
          metadata: { tool_progress: msg.tool_progress },
        });
      }
      break;

    case 'system_info':
      if (msg.system_info) {
        addMessage({
          id: msg.message_id || 'sys_' + Date.now(),
          role: 'system',
          content: msg.system_info.message,
          timestamp: msg.timestamp || Date.now(),
          msgType: 'system_info',
          isStreaming: false,
        });
      }
      break;

    case 'voice_interim':
      if (msg.text !== undefined) {
        useChatStore.getState().setVoiceInterimText(msg.text);
      }
      break;

    case 'voice_final':
      if (msg.text !== undefined) {
        useChatStore.getState().setVoiceFinalText(msg.text);
      }
      break;

    case 'pong':
      break;

    default:
      break;
  }
}
