backend/src/auth/two-factor/two-factor.service.ts

import {
  Injectable,
  BadRequestException,
  UnauthorizedException,
} from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { JwtService } from '@nestjs/jwt';
import * as speakeasy from 'speakeasy';
import * as QRCode from 'qrcode';
import * as argon2 from 'argon2';
import * as crypto from 'crypto';
import { PrismaService } from '../../prisma/prisma.service';

@Injectable()
export class TwoFactorService {
  private readonly encryptionKey: Buffer;
  private readonly algorithm = 'aes-256-gcm';

  constructor(
    private readonly prisma: PrismaService,
    private readonly configService: ConfigService,
    private readonly jwtService: JwtService,
  ) {
    // Use JWT secret as base for encryption key (derive a 32-byte key)
    const secret = this.configService.get<string>('jwt.secret') || 'default-secret';
    this.encryptionKey = crypto.scryptSync(secret, 'salt', 32);
  }

  // ==========================================
  // ENCRYPTION HELPERS
  // ==========================================
  private encrypt(text: string): string {
    const iv = crypto.randomBytes(16);
    const cipher = crypto.createCipheriv(this.algorithm, this.encryptionKey, iv);
    let encrypted = cipher.update(text, 'utf8', 'hex');
    encrypted += cipher.final('hex');
    const authTag = cipher.getAuthTag();
    return `${iv.toString('hex')}:${authTag.toString('hex')}:${encrypted}`;
  }

  private decrypt(encryptedData: string): string {
    const [ivHex, authTagHex, encrypted] = encryptedData.split(':');
    const iv = Buffer.from(ivHex, 'hex');
    const authTag = Buffer.from(authTagHex, 'hex');
    const decipher = crypto.createDecipheriv(this.algorithm, this.encryptionKey, iv);
    decipher.setAuthTag(authTag);
    let decrypted = decipher.update(encrypted, 'hex', 'utf8');
    decrypted += decipher.final('utf8');
    return decrypted;
  }

  // ==========================================
  // SETUP 2FA - Generate secret and QR code
  // ==========================================
  async setupTwoFactor(userId: string): Promise<{ secret: string; qrCode: string; manualEntry: string }> {
    const user = await this.prisma.user.findUnique({
      where: { id: userId },
    });

    if (!user) {
      throw new BadRequestException('User not found');
    }

    if (user.twoFactorEnabled) {
      throw new BadRequestException('Two-factor authentication is already enabled');
    }

    // Generate secret
    const secret = speakeasy.generateSecret({
      name: `Re-Quest:${user.email}`,
      issuer: 'Re-Quest',
      length: 32,
    });

    // Store encrypted secret temporarily (not yet verified)
    const encryptedSecret = this.encrypt(secret.base32);
    await this.prisma.user.update({
      where: { id: userId },
      data: { twoFactorSecret: encryptedSecret },
    });

    // Generate QR code
    const qrCode = await QRCode.toDataURL(secret.otpauth_url!);

    return {
      secret: secret.base32,
      qrCode,
      manualEntry: secret.base32,
    };
  }

  // ==========================================
  // ENABLE 2FA - Verify token and enable
  // ==========================================
  async enableTwoFactor(userId: string, token: string): Promise<{ backupCodes: string[] }> {
    const user = await this.prisma.user.findUnique({
      where: { id: userId },
    });

    if (!user || !user.twoFactorSecret) {
      throw new BadRequestException('Please setup two-factor authentication first');
    }

    if (user.twoFactorEnabled) {
      throw new BadRequestException('Two-factor authentication is already enabled');
    }

    // Decrypt and verify token
    const decryptedSecret = this.decrypt(user.twoFactorSecret);
    const isValid = speakeasy.totp.verify({
      secret: decryptedSecret,
      encoding: 'base32',
      token,
      window: 1, // Allow 1 step before/after for clock drift
    });

    if (!isValid) {
      throw new BadRequestException('Invalid verification code');
    }

    // Generate backup codes
    const backupCodes = await this.generateBackupCodesInternal();
    const hashedCodes = await Promise.all(
      backupCodes.map((code) => argon2.hash(code)),
    );

    // Enable 2FA
    await this.prisma.user.update({
      where: { id: userId },
      data: {
        twoFactorEnabled: true,
        twoFactorBackupCodes: JSON.stringify(hashedCodes),
        twoFactorVerifiedAt: new Date(),
      },
    });

    return { backupCodes };
  }

  // ==========================================
  // DISABLE 2FA
  // ==========================================
  async disableTwoFactor(userId: string, password: string): Promise<void> {
    const user = await this.prisma.user.findUnique({
      where: { id: userId },
    });

    if (!user) {
      throw new BadRequestException('User not found');
    }

    if (!user.twoFactorEnabled) {
      throw new BadRequestException('Two-factor authentication is not enabled');
    }

    // Verify password
    if (!user.password) {
      throw new BadRequestException('Cannot disable 2FA for social login accounts');
    }

    const isPasswordValid = await argon2.verify(user.password, password);
    if (!isPasswordValid) {
      throw new UnauthorizedException('Invalid password');
    }

    // Disable 2FA
    await this.prisma.user.update({
      where: { id: userId },
      data: {
        twoFactorEnabled: false,
        twoFactorSecret: null,
        twoFactorBackupCodes: null,
        twoFactorVerifiedAt: null,
      },
    });

    // Invalidate all sessions except current
    await this.prisma.session.deleteMany({
      where: { userId },
    });
  }

  // ==========================================
  // VERIFY TOKEN (during login)
  // ==========================================
  async verifyToken(userId: string, token: string): Promise<boolean> {
    const user = await this.prisma.user.findUnique({
      where: { id: userId },
    });

    if (!user || !user.twoFactorEnabled || !user.twoFactorSecret) {
      throw new BadRequestException('Two-factor authentication is not enabled');
    }

    const decryptedSecret = this.decrypt(user.twoFactorSecret);
    const isValid = speakeasy.totp.verify({
      secret: decryptedSecret,
      encoding: 'base32',
      token,
      window: 1,
    });

    return isValid;
  }

  // ==========================================
  // VERIFY BACKUP CODE
  // ==========================================
  async verifyBackupCode(userId: string, code: string): Promise<boolean> {
    const user = await this.prisma.user.findUnique({
      where: { id: userId },
    });

    if (!user || !user.twoFactorEnabled || !user.twoFactorBackupCodes) {
      throw new BadRequestException('Two-factor authentication is not enabled');
    }

    const hashedCodes: string[] = JSON.parse(user.twoFactorBackupCodes);

    // Find and verify the backup code
    for (let i = 0; i < hashedCodes.length; i++) {
      try {
        const isValid = await argon2.verify(hashedCodes[i], code);
        if (isValid) {
          // Remove used code
          hashedCodes.splice(i, 1);
          await this.prisma.user.update({
            where: { id: userId },
            data: {
              twoFactorBackupCodes: JSON.stringify(hashedCodes),
            },
          });
          return true;
        }
      } catch {
        continue;
      }
    }

    return false;
  }

  // ==========================================
  // REGENERATE BACKUP CODES
  // ==========================================
  async regenerateBackupCodes(userId: string, password: string): Promise<{ backupCodes: string[] }> {
    const user = await this.prisma.user.findUnique({
      where: { id: userId },
    });

    if (!user) {
      throw new BadRequestException('User not found');
    }

    if (!user.twoFactorEnabled) {
      throw new BadRequestException('Two-factor authentication is not enabled');
    }

    // Verify password
    if (!user.password) {
      throw new BadRequestException('Cannot regenerate codes for social login accounts');
    }

    const isPasswordValid = await argon2.verify(user.password, password);
    if (!isPasswordValid) {
      throw new UnauthorizedException('Invalid password');
    }

    // Generate new backup codes
    const backupCodes = await this.generateBackupCodesInternal();
    const hashedCodes = await Promise.all(
      backupCodes.map((code) => argon2.hash(code)),
    );

    await this.prisma.user.update({
      where: { id: userId },
      data: {
        twoFactorBackupCodes: JSON.stringify(hashedCodes),
      },
    });

    return { backupCodes };
  }

  // ==========================================
  // GET 2FA STATUS
  // ==========================================
  async getStatus(userId: string): Promise<{ enabled: boolean; verifiedAt: Date | null }> {
    const user = await this.prisma.user.findUnique({
      where: { id: userId },
      select: {
        twoFactorEnabled: true,
        twoFactorVerifiedAt: true,
      },
    });

    if (!user) {
      throw new BadRequestException('User not found');
    }

    return {
      enabled: user.twoFactorEnabled,
      verifiedAt: user.twoFactorVerifiedAt,
    };
  }

  // ==========================================
  // GENERATE TEMP TOKEN FOR 2FA FLOW
  // ==========================================
  async generateTempToken(userId: string): Promise<string> {
    const user = await this.prisma.user.findUnique({
      where: { id: userId },
      select: { email: true },
    });

    if (!user) {
      throw new BadRequestException('User not found');
    }

    return this.jwtService.sign(
      { sub: userId, email: user.email, type: '2fa-pending' },
      {
        secret: this.configService.get<string>('jwt.secret'),
        expiresIn: '5m',
      },
    );
  }

  // ==========================================
  // VERIFY TEMP TOKEN
  // ==========================================
  verifyTempToken(token: string): string {
    try {
      const payload = this.jwtService.verify(token, {
        secret: this.configService.get<string>('jwt.secret'),
      });

      if (payload.type !== '2fa-pending') {
        throw new UnauthorizedException('Invalid temporary token');
      }

      return payload.sub;
    } catch {
      throw new UnauthorizedException('Invalid or expired temporary token');
    }
  }

  // ==========================================
  // HELPER: Generate backup codes
  // ==========================================
  private async generateBackupCodesInternal(): Promise<string[]> {
    const codes: string[] = [];
    for (let i = 0; i < 8; i++) {
      // Generate 8 character alphanumeric codes
      const code = crypto.randomBytes(4).toString('hex').toUpperCase();
      codes.push(code);
    }
    return codes;
  }
}