import jwt from "jsonwebtoken";
import u from "@/utils";
import { Namespace, Socket } from "socket.io";
import * as agent from "@/agents/productionAgent/index";
import ResTool from "@/socket/resTool";

async function verifyToken(rawToken: string): Promise<Boolean> {
  const setting = await u.db("o_setting").where("key", "tokenKey").select("value").first();
  if (!setting) return false;
  const { value: tokenKey } = setting;
  if (!rawToken) return false;
  const token = rawToken.replace("Bearer ", "");
  try {
    jwt.verify(token, tokenKey as string);
    return true;
  } catch (err) {
    return false;
  }
}

export default (nsp: Namespace) => {
  nsp.on("connection", async (socket: Socket) => {
    const token = socket.handshake.auth.token;
    if (!token || !(await verifyToken(token))) {
      console.log("[productionAgent] 连接失败，token无效");
      socket.disconnect();
      return;
    }
    let isolationKey = socket.handshake.auth.isolationKey;
    if (!isolationKey) {
      console.log("[productionAgent] 连接失败，缺少 isolationKey");
      socket.disconnect();
      return;
    }

    console.log("[productionAgent] 已连接:", socket.id);

    const resTool = new ResTool(socket, {
      projectId: socket.handshake.auth.projectId,
      scriptId: socket.handshake.auth.scriptId,
    });
    let abortController: AbortController | null = null;

    socket.on("message", async (text: string) => {
      abortController?.abort();
      abortController = new AbortController();
      const currentController = abortController;
      console.log("%c Line:30 🍑 isolationKey", "background:#e41a6a", isolationKey);

      const textStream = await agent.decisionAI({ socket, isolationKey, text, abortSignal: currentController.signal, resTool });

      let msg = resTool.textMessage();

      try {
        for await (const chunk of textStream) {
          msg.send(chunk);
        }
      } catch (err: any) {
        if (err.name !== "AbortError") throw err;
      } finally {
        msg.end();
        if (abortController === currentController) {
          abortController = null;
        }
      }
    });
    socket.on("setModelData", async (data: any) => {
      resTool.data.imageModel = data;
    });
    socket.on("setKeyScript", async (data: any) => {
      isolationKey = data.key;
      resTool.data.scriptId = data.scriptId;
    });
    socket.on("stop", () => {
      abortController?.abort();
      abortController = null;
    });
  });
};
