import { jwtVerify, SignJWT, type JWTPayload } from "jose";
import { Role } from "@prisma/client";
import { env } from "@/server/env";

export type AccessClaims = JWTPayload & {
  sub: string;
  sessionId: string;
  roles: Role[];
  tokenType: "access";
};

export type RefreshClaims = JWTPayload & {
  sub: string;
  sessionId: string;
  tokenType: "refresh";
};

function key(secret: string) {
  return new TextEncoder().encode(secret);
}

export async function signAccessToken(input: { userId: string; sessionId: string; roles: Role[] }) {
  const config = env();
  return new SignJWT({
    sessionId: input.sessionId,
    roles: input.roles,
    tokenType: "access"
  })
    .setProtectedHeader({ alg: "HS256" })
    .setSubject(input.userId)
    .setIssuedAt()
    .setExpirationTime("15m")
    .sign(key(config.JWT_ACCESS_SECRET));
}

export async function signRefreshToken(input: { userId: string; sessionId: string }) {
  const config = env();
  return new SignJWT({
    sessionId: input.sessionId,
    tokenType: "refresh"
  })
    .setProtectedHeader({ alg: "HS256" })
    .setSubject(input.userId)
    .setIssuedAt()
    .setExpirationTime("30d")
    .sign(key(config.JWT_REFRESH_SECRET));
}

export async function verifyAccessToken(token: string) {
  const config = env();
  const { payload } = await jwtVerify(token, key(config.JWT_ACCESS_SECRET));

  if (payload.tokenType !== "access" || !payload.sub || !payload.sessionId || !Array.isArray(payload.roles)) {
    throw new Error("Invalid access token");
  }

  return payload as AccessClaims;
}

export async function verifyRefreshToken(token: string) {
  const config = env();
  const { payload } = await jwtVerify(token, key(config.JWT_REFRESH_SECRET));

  if (payload.tokenType !== "refresh" || !payload.sub || !payload.sessionId) {
    throw new Error("Invalid refresh token");
  }

  return payload as RefreshClaims;
}