import { WebSocketServer, WebSocket } from "ws";
import { randomBytes } from "crypto";
import type { IncomingMessage } from "http";
import { prisma } from "./prisma";
import { createEphemeralPreAuthKey, getHeadscaleUserId } from "./headscale";

interface NodeMessage {
  action: string;
  nodeId?: string;
  code?: string;
  instanceId?: string;
  type?: string;
  port?: number;
  composeConfig?: string;
  studentName?: string;
  error?: string;
  tailscaleIp?: string;
  token?: string;
}

const nodes = new Map<string, WebSocket>();

interface AttemptWindow {
  count: number;
  firstAttempt: number;
}

const activationAttemptsByCode = new Map<string, AttemptWindow>();
const activationAttemptsByNode = new Map<string, AttemptWindow>();
const MAX_ACTIVATION_ATTEMPTS = 5;
const ACTIVATION_WINDOW_MS = 15 * 60 * 1000;

const HEADSCALE_USER = "studioe5";
const HEADSCALE_AGENT_TAG = "tag:student-agent";
const HEADSCALE_KEY_EXPIRATION_MINUTES = 15;

let headscaleUserIdCache: string | null = null;

function recordActivationAttempt(map: Map<string, AttemptWindow>, key: string): boolean {
  const now = Date.now();
  const win = map.get(key);
  if (!win || now - win.firstAttempt > ACTIVATION_WINDOW_MS) {
    map.set(key, { count: 1, firstAttempt: now });
    return true;
  }
  win.count++;
  return win.count <= MAX_ACTIVATION_ATTEMPTS;
}

function clearActivationAttempts(code: string, nodeId: string) {
  activationAttemptsByCode.delete(code);
  activationAttemptsByNode.delete(nodeId);
}

function generateNodeToken(): string {
  return randomBytes(32).toString("hex");
}

function getBearerToken(req: IncomingMessage): string | null {
  const auth = req.headers.authorization || "";
  const match = auth.match(/^Bearer\s+(\S+)$/i);
  return match ? match[1] : null;
}

function close(ws: WebSocket, code: number, reason: string) {
  try {
    ws.close(code, reason);
  } catch {
    // ignore
  }
}

export function initWebSocketServer(wss: WebSocketServer) {
  wss.on("connection", (ws: WebSocket, req: IncomingMessage) => {
    let nodeId: string | null = null;
    let authenticated = false;
    const token = getBearerToken(req);

    console.log("[WS] New connection", token ? "(token provided)" : "(no token)");

    ws.on("message", async (raw) => {
      try {
        const msg: NodeMessage = JSON.parse(raw.toString());
        console.log("[WS] Received:", msg.action, "from", msg.nodeId || nodeId);

        if (msg.action === "register" && msg.nodeId) {
          const id = msg.nodeId;
          const existing = await prisma.node.findUnique({ where: { id } });

          if (token) {
            // Token supplied: it must match the stored token for this node.
            if (!existing || existing.token !== token) {
              console.log("[WS] Invalid token for node", id);
              close(ws, 1008, "invalid token");
              return;
            }
            authenticated = true;
          } else if (existing && existing.token) {
            // Existing node has a token but none was supplied.
            console.log("[WS] Missing token for node", id);
            close(ws, 1008, "missing token");
            return;
          } else if (existing) {
            // Migration path: existing node without a token gets one on first register.
            const newToken = generateNodeToken();
            await prisma.node.update({
              where: { id },
              data: { token: newToken, status: "online", lastSeen: new Date() },
            });
            ws.send(JSON.stringify({ action: "set_token", token: newToken }));
            authenticated = true;
          }
          // If the node does not exist yet, we stay unauthenticated until activation.

          nodeId = id;
          if (authenticated) {
            const existing = nodes.get(id);
            if (existing && existing !== ws && existing.readyState === WebSocket.OPEN) {
              console.log("[WS] Superseding previous connection for", id);
              existing.close(1008, "superseded");
            }
            nodes.set(id, ws);
            await prisma.node.upsert({
              where: { id },
              update: { status: "online", lastSeen: new Date() },
              create: { id, status: "online", lastSeen: new Date() },
            });
          }
          ws.send(JSON.stringify({ action: "registered" }));
          return;
        }

        if (msg.action === "activate" && msg.code && msg.nodeId) {
          const id = msg.nodeId;
          nodeId = id;

          if (!recordActivationAttempt(activationAttemptsByCode, msg.code) ||
              !recordActivationAttempt(activationAttemptsByNode, id)) {
            console.log("[WS] Too many activation attempts for code/node", msg.code, id);
            ws.send(JSON.stringify({ action: "activation_failed", error: "Too many attempts" }));
            return;
          }

          const existing = await prisma.node.findUnique({ where: { id } });
          if (existing && existing.token && (!authenticated || nodeId !== id)) {
            console.log("[WS] Node already activated and not authenticated:", id);
            ws.send(JSON.stringify({ action: "activation_failed", error: "Node already activated" }));
            return;
          }
          const student = await prisma.student.findUnique({
            where: { activationCode: msg.code },
          });
          if (!student) {
            console.log("[WS] Invalid code:", msg.code);
            ws.send(JSON.stringify({ action: "activation_failed", error: "Invalid code" }));
            return;
          }
          if (!student.activationCodeExpiresAt || student.activationCodeExpiresAt < new Date()) {
            console.log("[WS] Expired code:", msg.code);
            ws.send(JSON.stringify({ action: "activation_failed", error: "Code expired" }));
            return;
          }

          const newToken = generateNodeToken();
          await prisma.node.upsert({
            where: { id },
            update: {
              studentId: student.id,
              status: "online",
              lastSeen: new Date(),
              token: newToken,
            },
            create: {
              id,
              studentId: student.id,
              status: "online",
              lastSeen: new Date(),
              token: newToken,
            },
          });

          // Invalidate the activation code so it cannot be reused.
          await prisma.student.update({
            where: { id: student.id },
            data: { activationCode: null, activationCodeExpiresAt: null },
          });
          clearActivationAttempts(msg.code, id);

          authenticated = true;
          const previous = nodes.get(id);
          if (previous && previous !== ws && previous.readyState === WebSocket.OPEN) {
            console.log("[WS] Superseding previous connection for", id);
            previous.close(1008, "superseded");
          }
          nodes.set(id, ws);
          const headscaleUrl = process.env.HEADSCALE_URL;
          const headscaleApiKey = process.env.HEADSCALE_API_KEY;
          const reusableAuthKey = process.env.HEADSCALE_AUTH_KEY;

          if (!headscaleUrl) {
            console.log("[WS] HEADSCALE_URL missing");
            ws.send(JSON.stringify({ action: "activation_failed", error: "Server misconfiguration" }));
            return;
          }

          let headscaleAuthKey: string;
          try {
            if (headscaleApiKey) {
              if (!headscaleUserIdCache) {
                headscaleUserIdCache = await getHeadscaleUserId(headscaleUrl, headscaleApiKey, HEADSCALE_USER);
              }
              headscaleAuthKey = await createEphemeralPreAuthKey(headscaleUrl, headscaleApiKey, headscaleUserIdCache, {
                expirationMinutes: HEADSCALE_KEY_EXPIRATION_MINUTES,
                aclTags: [HEADSCALE_AGENT_TAG],
              });
              console.log("[WS] Generated ephemeral Headscale key for", id);
            } else if (reusableAuthKey) {
              console.log("[WS] HEADSCALE_API_KEY not set, falling back to reusable HEADSCALE_AUTH_KEY");
              headscaleAuthKey = reusableAuthKey;
            } else {
              console.log("[WS] No Headscale key available");
              ws.send(JSON.stringify({ action: "activation_failed", error: "Server misconfiguration" }));
              return;
            }
          } catch (err) {
            console.error("[WS] Failed to create ephemeral Headscale key:", err);
            ws.send(JSON.stringify({ action: "activation_failed", error: "Failed to create VPN key" }));
            return;
          }

          console.log("[WS] Activated:", student.firstName, student.lastName, "on", id);
          ws.send(JSON.stringify({
            action: "activated",
            studentId: student.id,
            studentName: `${student.firstName} ${student.lastName}`,
            headscaleUrl,
            headscaleAuthKey,
            token: newToken,
          }));
          return;
        }

        if (!authenticated || !nodeId) {
          console.log("[WS] Unauthenticated message", msg.action, "ignored");
          return;
        }

        if (msg.action === "heartbeat") {
          await prisma.node.upsert({
            where: { id: nodeId },
            update: { lastSeen: new Date() },
            create: { id: nodeId, status: "online", lastSeen: new Date() },
          });
          return;
        }

        if (msg.action === "tailscale_ip" && msg.tailscaleIp) {
          await prisma.node.update({
            where: { id: nodeId },
            data: { tailscaleIp: msg.tailscaleIp },
          });
          console.log("[WS] Tailscale IP updated for", nodeId, ":", msg.tailscaleIp);
          return;
        }

        if (msg.action === "instance_started" && msg.instanceId) {
          const { count } = await prisma.instance.updateMany({
            where: { id: msg.instanceId },
            data: { status: "running" },
          });
          if (count) console.log("[WS] Instance started:", msg.instanceId);
          return;
        }

        if (msg.action === "instance_stopped" && msg.instanceId) {
          const { count } = await prisma.instance.updateMany({
            where: { id: msg.instanceId },
            data: { status: "stopped" },
          });
          if (count) console.log("[WS] Instance stopped:", msg.instanceId);
          return;
        }

        if (msg.action === "instance_deleted" && msg.instanceId) {
          const { count } = await prisma.instance.deleteMany({
            where: { id: msg.instanceId },
          });
          if (count) console.log("[WS] Instance deleted:", msg.instanceId);
          return;
        }

        if (msg.action === "instance_error" && msg.instanceId) {
          const { count } = await prisma.instance.updateMany({
            where: { id: msg.instanceId },
            data: { status: "error" },
          });
          if (count) console.log("[WS] Instance error:", msg.instanceId);
          return;
        }
      } catch (err) {
        console.error("WS error:", err);
      }
    });

    ws.on("close", async () => {
      console.log("[WS] Connection closed for", nodeId);
      if (nodeId) {
        nodes.delete(nodeId);
        await prisma.node.upsert({
          where: { id: nodeId },
          update: { status: "offline" },
          create: { id: nodeId, status: "offline" },
        });
      }
    });
  });
}

export function sendToNode(nodeId: string, message: object) {
  const ws = nodes.get(nodeId);
  if (ws && ws.readyState === WebSocket.OPEN) {
    ws.send(JSON.stringify(message));
    return true;
  }
  return false;
}