import * as jose from 'jose';
import { env } from '../config/env';
import { fetchVaultKV2 } from '../config/vault';
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;
}

type KeyType = Awaited<ReturnType<typeof jose.importSPKI>>;

// Whitelist надёжных асимметричных алгоритмов. Никогда не разрешаем 'none'/HS*
// (HS — симметричные, могли бы быть подставлены через algorithm confusion).
const ALLOWED_ALGORITHMS = new Set(['RS256', 'RS384', 'RS512', 'ES256', 'ES384', 'ES512', 'EdDSA', 'PS256', 'PS384', 'PS512']);

if (!ALLOWED_ALGORITHMS.has(env.jwt.algorithm)) {
  throw new Error(`JWT_ALGORITHM "${env.jwt.algorithm}" not allowed. Use one of: ${[...ALLOWED_ALGORITHMS].join(', ')}`);
}

// Live key store — атомарно подменяется через swapKeyMap()
let keyMap: Map<string, KeyType> = new Map();

export function swapKeyMap(newMap: Map<string, KeyType>): void {
  keyMap = newMap;
}

export function getKeyMapSize(): number {
  return keyMap.size;
}

/**
 * Pre-fetch JWT public keys from Vault, не мутируя глобальный keyMap.
 * Возвращает новую Map для атомарного swap'а.
 */
export async function fetchJwtKeysFromVault(
  vaultAddr: string,
  vaultToken: string,
  mount: string,
  kidPath: string,
  kidsPrefix: string,
): Promise<Map<string, KeyType>> {
  const kidData = await fetchVaultKV2(vaultAddr, vaultToken, mount, kidPath);
  if (!kidData) {
    throw new Error('Failed to read JWT kid config from Vault');
  }

  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) {
    throw new Error('No active/previous kids found in Vault');
  }

  const next = new Map<string, KeyType>();

  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;
    }

    const key = await jose.importSPKI(kidSecret.public_key, env.jwt.algorithm);
    next.set(kid, key);
  }

  if (next.size === 0) {
    throw new Error('No public keys could be loaded from Vault');
  }

  return next;
}

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

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

    if (!kid || typeof kid !== 'string' || !/^[A-Za-z0-9_-]{1,64}$/.test(kid)) {
      throw Object.assign(new Error('Missing or invalid kid in token header'), { status: 401 });
    }

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

    // Двойная защита от algorithm confusion: проверяем точное совпадение
    if (header.alg !== env.jwt.algorithm || !ALLOWED_ALGORITHMS.has(String(header.alg))) {
      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 });
  }

  // Строгая валидация sub/sid — иначе number/__proto__/10MB строки попадают в PG / в req.auth.
  const SUB_RE = /^[A-Za-z0-9_-]{1,64}$/;
  if (typeof payload.sub !== 'string' || !SUB_RE.test(payload.sub)) {
    throw Object.assign(new Error('Invalid sub claim'), { status: 401 });
  }
  if (typeof payload.sid !== 'string' || !SUB_RE.test(payload.sid)) {
    throw Object.assign(new Error('Invalid sid claim'), { 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,
    },
  };
}