// 设置
import { useLoadingStore } from "@/stores/loadingStore";
import { io, Socket as SocketIOClient } from "socket.io-client";

// ─── 会话加密接口 ───────────────────────────────────────────────
export interface SessionCrypto {
  aesKey: CryptoKey; // AES-128-GCM，不可导出
}

// ─── AES-GCM 加密 / 解密 ───────────────────────────────────────
async function encryptPayload(plain: string, aesKey: CryptoKey): Promise<string> {
  const iv = crypto.getRandomValues(new Uint8Array(12));
  const encoded = new TextEncoder().encode(plain);
  const cipher = await crypto.subtle.encrypt({ name: "AES-GCM", iv }, aesKey, encoded);
  const out = new Uint8Array(iv.byteLength + cipher.byteLength);
  out.set(iv, 0);
  out.set(new Uint8Array(cipher), iv.byteLength);
  let binary = "";
  for (let i = 0; i < out.length; i++) binary += String.fromCharCode(out[i]);
  return JSON.stringify({ data: btoa(binary) });
}

async function decryptPayload(raw: unknown, aesKey: CryptoKey): Promise<string> {
  const rawStr =
    typeof raw === "string" ? raw :
    raw && typeof raw === "object" ? JSON.stringify(raw) : String(raw ?? "");

  let envelope: { data?: string };
  try { envelope = JSON.parse(rawStr); } catch { return rawStr; }
  if (!envelope?.data) return rawStr;

  const bytes = Uint8Array.from(atob(envelope.data), c => c.charCodeAt(0));
  const iv = bytes.slice(0, 12);
  const cipher = bytes.slice(12);
  try {
    const plain = await crypto.subtle.decrypt({ name: "AES-GCM", iv }, aesKey, cipher);
    return new TextDecoder().decode(plain);
  } catch {
    return rawStr;
  }
}

// ─── ECDH 密钥协商工具 ─────────────────────────────────────────

/**
 * 生成 P-256 临时密钥对，返回 { keyPair, clientPublicKeyB64 }
 */
export async function generateECDHKeyPair(): Promise<{
  keyPair: CryptoKeyPair;
  clientPublicKeyB64: string;
}> {
  const keyPair = await crypto.subtle.generateKey(
    { name: "ECDH", namedCurve: "P-256" },
    true,
    ["deriveBits"]
  );
  const pubKeyRaw = await crypto.subtle.exportKey("raw", keyPair.publicKey);
  const clientPublicKeyB64 = btoa(
    Array.from(new Uint8Array(pubKeyRaw)).map(b => String.fromCharCode(b)).join("")
  );
  return { keyPair, clientPublicKeyB64 };
}

/**
 * 用服务端公钥（base64 raw P-256）与给定的客户端私钥推导 AES-128-GCM 会话密钥
 */
export async function deriveSessionKey(
  serverPublicKeyB64: string,
  clientPrivateKey: CryptoKey
): Promise<SessionCrypto> {
  const serverPubKeyBytes = Uint8Array.from(atob(serverPublicKeyB64), c => c.charCodeAt(0));
  const serverPublicKey = await crypto.subtle.importKey(
    "raw", serverPubKeyBytes,
    { name: "ECDH", namedCurve: "P-256" }, false, []
  );
  const sharedBits = await crypto.subtle.deriveBits(
    { name: "ECDH", public: serverPublicKey }, clientPrivateKey, 256
  );
  const hkdfKey = await crypto.subtle.importKey("raw", sharedBits, "HKDF", false, ["deriveKey"]);
  const aesKey = await crypto.subtle.deriveKey(
    {
      name: "HKDF", hash: "SHA-256",
      salt: new Uint8Array(32),
      info: new TextEncoder().encode("socket-aes-key"),
    },
    hkdfKey,
    { name: "AES-GCM", length: 128 },
    false,
    ["encrypt", "decrypt"]
  );
  return { aesKey };
}

/** 断线/握手阶段队列最大长度，防止内存无限增长 */
const MAX_QUEUE_SIZE = 200;

class Socket {
  url: string;
  private token: string;
  private sessionCrypto: SessionCrypto | null;
  private ecdhKeyPair: CryptoKeyPair | null = null;
  private clientPublicKeyB64: string | null = null;
  /** 握手全部完成（ECDH + login）后才为 true，期间消息也入队 */
  private isReady = false;
  socket: SocketIOClient | null = null;
  listeners: { [key: string]: Function[] } = {};
  private messageQueue: any[] = []; // 断连/握手期间暂存消息的队列

  constructor(url: string, token = "", sessionCrypto: SessionCrypto | null = null) {
    this.url = url;
    this.token = token;
    this.sessionCrypto = sessionCrypto;
    this.init();
    this.setupVisibilityListener();
  }

  /** 懒初始化 ECDH 密钥对（只生成一次，重连时复用） */
  private async initECDH() {
    if (!this.ecdhKeyPair) {
      const { keyPair, clientPublicKeyB64 } = await generateECDHKeyPair();
      this.ecdhKeyPair = keyPair;
      this.clientPublicKeyB64 = clientPublicKeyB64;
    }
  }

  /**
   * 通过 Socket.IO key_exchange 事件与服务端协商会话密钥。
   * 每次 connect（包括服务端重启后重连）都调用，确保密钥始终有效。
   */
  private performKeyExchange(): Promise<void> {
    return new Promise<void>((resolve) => {
      // 3 秒超时：若服务端不响应则无加密继续
      const timeout = setTimeout(() => {
        this.socket?.off('key_exchange_result', onResult);
        this.sessionCrypto = null;
        resolve();
      }, 3000);

      const onResult = async (serverPubKeyB64: string) => {
        clearTimeout(timeout);
        try {
          this.sessionCrypto = await deriveSessionKey(serverPubKeyB64, this.ecdhKeyPair!.privateKey);
        } catch (e) {
          console.error('[Socket] key derivation failed:', e);
          this.sessionCrypto = null;
        }
        resolve();
      };

      this.socket?.once('key_exchange_result', onResult);
      this.socket?.emit('key_exchange', this.clientPublicKeyB64);
    });
  }

  /**
   * 发送 login 并等待服务端回 {event:"login",content:"success"}。
   * 服务端发送 login success 时 client.State 已同步设置完毕，
   * 之后再冲刷队列才能保证消息不被 "State == nil" 守卫丢弃。
   */
  private async sendLoginAndWait(): Promise<void> {
    await new Promise<void>((resolve) => {
      let settled = false;
      const settle = () => {
        if (settled) return;
        settled = true;
        this.socket?.off('message', onRawMessage);
        clearTimeout(timer);
        resolve();
      };
      // 3 秒兜底：即使没收到确认也继续
      const timer = setTimeout(settle, 3000);

      const onRawMessage = async (raw: unknown) => {
        try {
          let text: string;
          if (this.sessionCrypto) {
            text = await decryptPayload(raw, this.sessionCrypto.aesKey);
          } else {
            text = typeof raw === 'string' ? raw : JSON.stringify(raw);
          }
          const parsed = JSON.parse(text);
          if (parsed?.event === 'login') settle();
        } catch { /* 忽略解析失败 */ }
      };

      // 先注册监听，再发 login，避免极速响应漏掉
      this.socket?.on('message', onRawMessage);
      this.sendRaw(JSON.stringify({ event: 'login', content: { token: this.token }, timestamp: Date.now() })).catch(() => settle());
    });
  }

  init() {
    if (this.socket) {
      return;
    }
    console.log("Socket initialized with URL:", this.url);
    this.socket = io(this.url, {
      path: "/socket.io",
      query: this.token ? { token: this.token } : undefined,
      reconnectionDelay: 1500,
      reconnectionAttempts: Infinity, // 服务端重启后持续重连，不放弃
    });

    // 连接事件处理（含重连）：每次都重新做 ECDH，解决服务端重启后密钥失效问题
    this.socket.on('connect', async () => {
      this.isReady = false; // 握手期间暂停直接发送，新消息继续入队
      // 清理上一次连接残留的 key_exchange_result 监听器，避免多次重连后堆积
      this.socket?.off('key_exchange_result');
      try {
        await this.initECDH();
        await this.performKeyExchange();
        // 等待服务端 login success 确认后再冲刷队列
        // 保证 client.State 已在服务端设置，避免消息被 "State==nil" 守卫丢弃
        await this.sendLoginAndWait();
      } catch (e) {
        console.error('[Socket] 握手阶段异常，将以无加密方式继续:', e);
        this.sessionCrypto = null;
      } finally {
        // 无论握手是否成功，都必须就绪并冲刷队列，避免消息永久滞留
        this.isReady = true;
        await this.flushMessageQueue(); // 连接后按序发送排队消息
        this.emit('open', { type: 'open' });
      }
    });

    // 消息接收（支持 AES-GCM 解密）
    this.socket.on('message', async (data) => {
      let plainText: string;
      if (this.sessionCrypto) {
        plainText = await decryptPayload(data, this.sessionCrypto.aesKey);
      } else {
        plainText = typeof data === 'string' ? data : JSON.stringify(data);
      }
      this.emit('message', plainText);
    });

    // 连接错误
    this.socket.on('connect_error', (error) => {
      this.emit('error', error);
    });

    // 断开连接
    this.socket.on('disconnect', (reason) => {
      this.isReady = false; // 断开后消息重新入队
      this.emit('close', { reason });
    });

    // 重连尝试
    this.socket.on('reconnect_attempt', (attemptNumber) => {
      this.emit('reconnect_attempt', attemptNumber);
    });

    // 重连成功
    this.socket.on('reconnect', (attemptNumber) => {
      this.emit('reconnect', attemptNumber);
    });

    // 重连失败
    this.socket.on('reconnect_failed', () => {
      useLoadingStore().setLoading(false); // 重连失败时关闭加载状态
      this.emit('reconnect_failed', { type: 'reconnect_failed' });
    });

    // 处理所有其他事件
    this.socket.onAny((eventName, ...args) => {
      if (!['connect', 'disconnect', 'error', 'reconnect_attempt',
        'reconnect', 'reconnect_failed', 'message'].includes(eventName)) {
        this.emit(eventName, args);
      }
    });
  }

  isConnected(): boolean {
    return this.socket?.connected ?? false;
  }

  /** 用于握手阶段的 login 事件，同样走加密通道 */
  private async sendRaw(data: string) {
    if (this.sessionCrypto) {
      const encrypted = await encryptPayload(data, this.sessionCrypto.aesKey);
      this.socket?.emit('message', encrypted);
    } else {
      this.socket?.emit('message', data);
    }
  }

  async send(data: string) {
    try {
      const payload = JSON.parse(data);

      // 添加时间戳
      const messageData = {
        ...payload,
        timestamp: payload.timestamp || Date.now()
      };

      // 未就绪（断连中或 ECDH/login 握手中）时统一入队，保证消息不丢失且顺序正确
      if (!this.isReady) {
        if (this.messageQueue.length < MAX_QUEUE_SIZE) {
          this.messageQueue.push(messageData);
        } else {
          console.warn('[Socket] 消息队列已满，丢弃消息:', messageData.event);
        }
        this.reconnectIfNeeded();
        return;
      }

      const serialized = JSON.stringify(messageData);
      if (this.sessionCrypto) {
        const encrypted = await encryptPayload(serialized, this.sessionCrypto.aesKey);
        this.socket?.emit('message', encrypted);
      } else {
        this.socket?.emit('message', serialized);
      }
    } catch (error) {
      console.error('Invalid message format. Must be a valid JSON string.', error);
    }
  }

  /** 按顺序逐条发送积压消息，保证 FIFO 且不会因并发导致乱序 */
  async flushMessageQueue() {
    if (this.messageQueue.length === 0) return;
    const queue = this.messageQueue.splice(0); // 原子取出，避免发送期间新消息混入
    for (const msg of queue) {
      if (!this.isReady || !this.socket?.connected) {
        // 发送途中再次断开，将剩余消息放回队首
        this.messageQueue.unshift(...queue.slice(queue.indexOf(msg)));
        break;
      }
      try {
        const serialized = JSON.stringify(msg);
        if (this.sessionCrypto) {
          const encrypted = await encryptPayload(serialized, this.sessionCrypto.aesKey);
          this.socket?.emit('message', encrypted);
        } else {
          this.socket?.emit('message', serialized);
        }
      } catch (e) {
        console.error('[Socket] flushMessageQueue 发送失败:', e);
        // 发送失败也放回队首
        this.messageQueue.unshift(...queue.slice(queue.indexOf(msg)));
        break;
      }
    }
  }

  reconnectIfNeeded() {
    if (!this.isConnected() && this.socket) {
      this.socket.connect();
    }
  }

  on(event: string, callback: Function) {
    // 需要经过本层中间件（如解密）的事件，统一走 this.listeners
    if (['open', 'close', 'error', 'reconnect', 'reconnect_attempt', 'reconnect_failed', 'message'].includes(event)) {
      if (!this.listeners[event]) {
        this.listeners[event] = [];
      }
      this.listeners[event].push(callback);
    } else {
      // 其他 Socket.IO 原生事件
      this.socket?.on(event, (...args) => callback(...args));
    }
  }

  off(event: string) {
    if (this.listeners[event]) {
      delete this.listeners[event];
    }
    this.socket?.off(event);
  }

  emit(event: string, data: any) {
    if (this.listeners[event]) {
      this.listeners[event].forEach(callback => callback(data));
    }
  }

  private handleVisibilityChange = () => {
    if (document.visibilityState === "visible" && !this.isConnected() && this.socket) {
      this.socket.connect();
    }
  };

  setupVisibilityListener() {
    document.addEventListener("visibilitychange", this.handleVisibilityChange);
  }

  disconnect() {
    document.removeEventListener("visibilitychange", this.handleVisibilityChange);
    this.socket?.disconnect();
  }
}

function useSocketIo(url: string, token = "", sessionCrypto: SessionCrypto | null = null) {
  const socket = new Socket(url, token, sessionCrypto);

  return {
    socket,
    send: socket.send.bind(socket),
    on: socket.on.bind(socket),
    off: socket.off.bind(socket),
    disconnect: socket.disconnect.bind(socket),
  };
}

export { useSocketIo, Socket };