import * as jose from 'jose';
import { env } from '../config/env';
import { logger } from '../lib/logger';

export interface AccessTokenPayload {
  sub: string;
  type: string;
  sid: string;
  iat: number;
  nbf: number;
  exp: number;
  iss?: string;
  aud?: string;
}

export interface AuthContext {
  userId: string;
  sid: string;
  token: AccessTokenPayload;
}

const keyMap = new Map<string, Awaited<ReturnType<typeof jose.importSPKI>>>();

export async function loadJwtKeysFromVault(
  vaultAddr: string,
  vaultToken: string,
  mount: string,
  kidPath: string,
  kidsPrefix: string,
): Promise<void> {
  const { fetchVaultKV2 } = await import('../config/vault');

  const kidData = await fetchVaultKV2(vaultAddr, vaultToken, mount, kidPath);
  if (!kidData) {
    logger.warn('Failed to read JWT kid config from Vault');
    return;
  }

  const kids: string[] = [];
  if (kidData.active) kids.push(kidData.active);
  if (kidData.previous && kidData.previous !== kidData.active) kids.push(kidData.previous);

  if (kids.length === 0) {
    logger.warn('No active/previous kids found in Vault');
    return;
  }

  for (const kid of kids) {
    const kidSecret = await fetchVaultKV2(vaultAddr, vaultToken, mount, `${kidsPrefix}/${kid}`);
    if (!kidSecret?.public_key) {
      logger.warn(`No public_key found for kid=${kid}`);
      continue;
    }

    try {
      const key = await jose.importSPKI(kidSecret.public_key, env.jwt.algorithm);
      keyMap.set(kid, key);
      logger.info(`Loaded JWT public key for kid=${kid}`);
    } catch (err: any) {
      logger.error(`Failed to import public key for kid=${kid}: ${err.message}`);
    }
  }

  logger.info(`JWT key store loaded: ${keyMap.size} key(s)`);
}

export async function verifyAccessToken(token: string): Promise<AuthContext> {
  let payload: jose.JWTPayload;

  try {
    const header = jose.decodeProtectedHeader(token);
    const kid = header.kid;

    if (!kid) {
      throw Object.assign(new Error('Missing kid in token header'), { status: 401 });
    }

    const key = keyMap.get(kid);
    if (!key) {
      logger.warn(`Unknown kid=${kid}`);
      throw Object.assign(new Error('Unknown token kid'), { status: 401 });
    }

    if (header.alg !== env.jwt.algorithm) {
      throw Object.assign(new Error('Invalid token algorithm'), { status: 401 });
    }

    const verifyOptions: jose.JWTVerifyOptions = {
      algorithms: [env.jwt.algorithm],
      clockTolerance: 10,
    };
    if (env.jwt.issuer) verifyOptions.issuer = env.jwt.issuer;
    if (env.jwt.audience) verifyOptions.audience = env.jwt.audience;

    const result = await jose.jwtVerify(token, key, verifyOptions);
    payload = result.payload;
  } catch (err: any) {
    if (err.status === 401) throw err;
    if (err.code === 'ERR_JWT_EXPIRED') {
      throw Object.assign(new Error('Token expired'), { status: 401 });
    }
    throw Object.assign(new Error('Invalid token'), { status: 401 });
  }

  if (payload.type !== 'access') {
    throw Object.assign(new Error('Invalid token type'), { status: 401 });
  }

  if (!payload.sub || !payload.sid) {
    throw Object.assign(new Error('Missing token claims'), { status: 401 });
  }

  return {
    userId: payload.sub,
    sid: payload.sid as string,
    token: {
      sub: payload.sub,
      type: payload.type as string,
      sid: payload.sid as string,
      iat: payload.iat!,
      nbf: payload.nbf!,
      exp: payload.exp!,
      iss: payload.iss,
      aud: typeof payload.aud === 'string' ? payload.aud : undefined,
    },
  };
}