import { createHash, createHmac, timingSafeEqual } from "node:crypto";

type TokenKind = "client_session" | "client_magic_link";
type MagicLinkPurpose = "invite" | "login";

type BaseClaims = {
  email: string;
  exp: number;
  iat: number;
  sub: string;
  type: TokenKind;
};

type ClientSessionClaims = BaseClaims & {
  type: "client_session";
};

type ClientMagicLinkClaims = BaseClaims & {
  jti: string;
  purpose: MagicLinkPurpose;
  redirect_to?: string;
  type: "client_magic_link";
};

type SignSessionOptions = {
  expiresInDays: number;
  secret: string;
};

type SignMagicLinkOptions = {
  expiresInMinutes: number;
  secret: string;
};

type VerifyOptions = {
  secret: string;
};

type BuildCookieOptions = {
  cookieDomain?: string;
  cookieName: string;
  maxAgeSeconds: number;
};

type SignClientSessionInput = {
  clientId: string;
  email: string;
};

type SignClientMagicLinkInput = {
  clientId: string;
  email: string;
  jti: string;
  purpose: MagicLinkPurpose;
  redirectTo?: string;
};

function encodeSegment(value: unknown) {
  return Buffer.from(JSON.stringify(value)).toString("base64url");
}

function decodeSegment<T>(segment: string): T | null {
  try {
    return JSON.parse(Buffer.from(segment, "base64url").toString("utf8")) as T;
  } catch {
    return null;
  }
}

function signToken(claims: ClientSessionClaims | ClientMagicLinkClaims, secret: string) {
  const header = encodeSegment({ alg: "HS256", typ: "JWT" });
  const payload = encodeSegment(claims);
  const signature = createHmac("sha256", secret).update(`${header}.${payload}`).digest("base64url");

  return `${header}.${payload}.${signature}`;
}

function verifyToken<T extends ClientSessionClaims | ClientMagicLinkClaims>(
  token: string,
  secret: string,
  expectedType: TokenKind
) {
  const segments = token.split(".");
  if (segments.length !== 3) {
    throw new Error("Invalid client session token");
  }

  const [header, payload, signature] = segments;
  const expectedSignature = createHmac("sha256", secret).update(`${header}.${payload}`).digest();
  const receivedSignature = Buffer.from(signature, "base64url");

  if (
    expectedSignature.length !== receivedSignature.length ||
    !timingSafeEqual(expectedSignature, receivedSignature)
  ) {
    throw new Error("Invalid client session token");
  }

  const claims = decodeSegment<T>(payload);
  if (!claims || claims.type !== expectedType) {
    throw new Error("Invalid client session token");
  }

  if (claims.exp <= Math.floor(Date.now() / 1000)) {
    throw new Error("Client session token expired");
  }

  return claims;
}

export function signClientSession(input: SignClientSessionInput, options: SignSessionOptions) {
  const now = Math.floor(Date.now() / 1000);
  return signToken(
    {
      email: input.email,
      exp: now + options.expiresInDays * 24 * 60 * 60,
      iat: now,
      sub: input.clientId,
      type: "client_session",
    },
    options.secret
  );
}

export function verifyClientSession(token: string, options: VerifyOptions) {
  return verifyToken<ClientSessionClaims>(token, options.secret, "client_session");
}

export function signClientMagicLink(input: SignClientMagicLinkInput, options: SignMagicLinkOptions) {
  const now = Math.floor(Date.now() / 1000);
  return signToken(
    {
      email: input.email,
      exp: now + options.expiresInMinutes * 60,
      iat: now,
      jti: input.jti,
      purpose: input.purpose,
      redirect_to: input.redirectTo,
      sub: input.clientId,
      type: "client_magic_link",
    },
    options.secret
  );
}

export function verifyClientMagicLink(token: string, options: VerifyOptions) {
  return verifyToken<ClientMagicLinkClaims>(token, options.secret, "client_magic_link");
}

export function buildClientSessionCookie(token: string, options: BuildCookieOptions) {
  const domainPart = options.cookieDomain ? `; Domain=${options.cookieDomain}` : "";
  return `${options.cookieName}=${token}; Path=/; HttpOnly; Secure; SameSite=Lax; Max-Age=${options.maxAgeSeconds}${domainPart}`;
}

export function clearClientSessionCookie(options: {
  cookieDomain?: string;
  cookieName: string;
}) {
  const domainPart = options.cookieDomain ? `; Domain=${options.cookieDomain}` : "";
  return `${options.cookieName}=; Path=/; HttpOnly; Secure; SameSite=Lax; Max-Age=0${domainPart}`;
}

export function readClientSessionCookie(cookieHeader: string | null | undefined, options: {
  cookieName: string;
}) {
  if (!cookieHeader) {
    return null;
  }

  const cookieMatch = cookieHeader.match(
    new RegExp(`(?:^|;\\s*)${options.cookieName}=([^;]+)`)
  );

  return cookieMatch?.[1] ?? null;
}

export function hashClientMagicLinkToken(token: string) {
  return createHash("sha256").update(token).digest("hex");
}

export type { ClientMagicLinkClaims, ClientSessionClaims, MagicLinkPurpose };
