import test from "node:test";
import assert from "node:assert/strict";
import { randomUUID } from "node:crypto";
import { WebSocket } from "ws";
import { createHttpServer } from "../src/app/http.ts";
import { attachTaskWebSocket } from "../src/app/ws.ts";
import { createStore } from "../src/store.ts";
import { hashRefreshToken, issueAccessToken } from "../src/shared/auth.ts";

test("task websocket authenticates and replays a subscribed task cursor", async () => {
  const secret = "test";
  const store = createStore();
  const server = createHttpServer(store, { host: "127.0.0.1", port: 0, accessTokenSecret: secret, accessTokenTtlSeconds: 900, refreshCookieName: "refresh", adminEmail: "admin@example.com", adminPassword: "change-me-now", adminMfaRequired: true });
  attachTaskWebSocket(server, store, secret, "refresh");
  await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
  const address = server.address();
  assert.ok(address && typeof address !== "string");
  const taskId = "ws-task";
  store.users.set("ws-user", { id: "ws-user", version: 1, password: "", verified: true, status: "active", mfaRequired: false, mfaEnabled: false, roles: ["user"], failedLoginCount: 0, createdAt: new Date().toISOString() });
  store.tasks.set(taskId, { id: taskId, ownerId: "ws-user", taskType: "image", modelProductId: "basic-image", status: "queued", estimatedCost: 1, reservedCost: 0, createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(), eventSequence: 0, routeSnapshot: { groupId: "image-default", groupVersion: 1, publicModelId: "basic-image-v1", routeVersion: 1, orderedCandidates: [], strategy: "strict", totalRetryBudget: 0 }, modelSnapshot: { publicModelId: "basic-image-v1", modelVersion: 1, name: "基础模型", tier: "basic", capabilities: ["image"] }, pricingSnapshot: { ruleId: "price-basic-image", ruleVersion: 1, publicModelId: "basic-image-v1", unitPrice: 1, quantity: 1, billingUnit: "output", multiplier: 1, balanceUnitVersion: 1 }, planSnapshot: { queuePriority: 0, maxConcurrent: 1 }, attempts: [], outputs: [] });
  const token = issueAccessToken("ws-user", "user", secret, 900);
  const socket = new WebSocket(`ws://127.0.0.1:${address.port}/api/v1/ws/tasks?access_token=${encodeURIComponent(token)}`);
  const messages: Array<{ type: string; event?: { taskId: string } }> = [];
  socket.on("message", (raw) => messages.push(JSON.parse(raw.toString())));
  await new Promise<void>((resolve) => socket.once("open", () => resolve()));
  socket.send(JSON.stringify({ type: "subscribe", taskId, cursor: 0 }));
  await new Promise((resolve) => setTimeout(resolve, 25));
  assert.equal(messages[0]?.type, "ready");
  store.events.set(taskId, []);
  const { appendEvent } = await import("../src/store.ts");
  appendEvent(store, taskId, { taskId, type: "task.progress", payload: { progress: 20 } });
  await new Promise((resolve) => setTimeout(resolve, 25));
  assert.ok(messages.some((message) => message.type === "event" && message.event?.taskId === taskId));
  socket.close();
  await new Promise<void>((resolve, reject) => server.close((error) => (error ? reject(error) : resolve())));
});

test("task websocket rejects a revoked refresh session family", async () => {
  const secret = "test";
  const store = createStore();
  const server = createHttpServer(store, { host: "127.0.0.1", port: 0, accessTokenSecret: secret, accessTokenTtlSeconds: 900, refreshCookieName: "refresh", adminEmail: "admin@example.com", adminPassword: "change-me-now", adminMfaRequired: true });
  attachTaskWebSocket(server, store, secret, "refresh");
  await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
  const address = server.address(); assert.ok(address && typeof address !== "string");
  const userId = "ws-cookie-user"; const familyId = randomUUID(); const refresh = randomUUID();
  store.users.set(userId, { id: userId, version: 1, password: "", verified: true, status: "active", mfaRequired: false, mfaEnabled: false, roles: ["user"], failedLoginCount: 0, createdAt: new Date().toISOString() });
  store.sessionFamilies.set(familyId, { userId, scope: "user" });
  store.sessions.set(hashRefreshToken(refresh), { userId, scope: "user", expiresAt: Date.now() + 60_000, absoluteExpiresAt: Date.now() + 60_000, createdAt: Date.now(), rotationId: randomUUID(), familyId });
  const first = new WebSocket(`ws://127.0.0.1:${address.port}/api/v1/ws/tasks`, { headers: { Cookie: `refresh=${refresh}` } });
  await new Promise<void>((resolve, reject) => { first.once("open", () => resolve()); first.once("error", reject); });
  store.sessionFamilies.get(familyId)!.revokedAt = Date.now();
  first.send(JSON.stringify({ type: "ping" }));
  await Promise.race([
    new Promise<void>((resolve) => first.once("close", () => resolve())),
    new Promise<void>((resolve) => setTimeout(resolve, 1_000)),
  ]);
  assert.equal(first.readyState, WebSocket.CLOSED);
  const second = new WebSocket(`ws://127.0.0.1:${address.port}/api/v1/ws/tasks`, { headers: { Cookie: `refresh=${refresh}` } });
  await Promise.race([
    new Promise<void>((resolve) => { second.once("close", () => resolve()); second.once("error", () => resolve()); }),
    new Promise<void>((resolve) => setTimeout(resolve, 1_000)),
  ]);
  second.terminate();
  await new Promise<void>((resolve, reject) => server.close((error) => error ? reject(error) : resolve()));
});