import { randomBytes } from 'node:crypto';
import { ForbiddenException, Injectable, NestMiddleware } from '@nestjs/common';
import { NextFunction, Request, Response } from 'express';

const CSRF_COOKIE = 'XSRF-TOKEN';
const CSRF_HEADER = 'x-csrf-token';
const TOKEN_LENGTH = 32;

const SAFE_METHODS = new Set(['GET', 'HEAD', 'OPTIONS']);

@Injectable()
export class CsrfMiddleware implements NestMiddleware {
  use(req: Request, res: Response, next: NextFunction): void {
    // Safe methods: ensure a CSRF cookie exists for the client to read
    if (SAFE_METHODS.has(req.method)) {
      this.ensureCsrfCookie(req, res);
      return next();
    }

    // State-changing methods: validate the double-submit token
    const cookieToken = req.cookies?.[CSRF_COOKIE] as string | undefined;
    const headerToken = req.headers[CSRF_HEADER] as string | undefined;

    if (!cookieToken || !headerToken || cookieToken !== headerToken) {
      throw new ForbiddenException('CSRF token missing or invalid');
    }

    // Rotate token after successful validation
    this.setCsrfCookie(res);
    next();
  }

  private ensureCsrfCookie(req: Request, res: Response): void {
    if (!req.cookies?.[CSRF_COOKIE]) {
      this.setCsrfCookie(res);
    }
  }

  private setCsrfCookie(res: Response): void {
    const token = randomBytes(TOKEN_LENGTH).toString('hex');
    const isProduction = process.env['NODE_ENV'] === 'production';
    res.cookie(CSRF_COOKIE, token, {
      httpOnly: false, // Frontend must read this cookie
      secure: isProduction,
      sameSite: isProduction ? 'strict' : 'lax',
      path: '/',
    });
  }
}