wraith/backend/src/auth/auth.service.ts

import { Injectable, UnauthorizedException, BadRequestException, NotFoundException, ForbiddenException } from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import { PrismaService } from '../prisma/prisma.service';
import * as bcrypt from 'bcrypt';
import * as speakeasy from 'speakeasy';
import * as QRCode from 'qrcode';

@Injectable()
export class AuthService {
  constructor(
    private prisma: PrismaService,
    private jwt: JwtService,
  ) {}

  async login(email: string, password: string, totpCode?: string) {
    const user = await this.prisma.user.findUnique({ where: { email } });
    if (!user) throw new UnauthorizedException('Invalid credentials');

    const valid = await bcrypt.compare(password, user.passwordHash);
    if (!valid) throw new UnauthorizedException('Invalid credentials');

    // If TOTP is enabled, require a valid code
    if (user.totpEnabled && user.totpSecret) {
      if (!totpCode) {
        // Signal frontend to show TOTP input
        return { requires_totp: true };
      }
      const verified = speakeasy.totp.verify({
        secret: user.totpSecret,
        encoding: 'base32',
        token: totpCode,
        window: 1,
      });
      if (!verified) throw new UnauthorizedException('Invalid TOTP code');
    }

    const payload = { sub: user.id, email: user.email, role: user.role };
    return {
      access_token: this.jwt.sign(payload),
      user: { id: user.id, email: user.email, displayName: user.displayName, role: user.role },
    };
  }

  async getProfile(userId: number) {
    const user = await this.prisma.user.findUnique({ where: { id: userId } });
    if (!user) throw new UnauthorizedException();
    return {
      id: user.id,
      email: user.email,
      displayName: user.displayName,
      role: user.role,
      totpEnabled: user.totpEnabled,
    };
  }

  async updateProfile(userId: number, data: { email?: string; displayName?: string; currentPassword?: string; newPassword?: string }) {
    const user = await this.prisma.user.findUnique({ where: { id: userId } });
    if (!user) throw new UnauthorizedException();

    const update: any = {};

    if (data.email && data.email !== user.email) {
      update.email = data.email;
    }
    if (data.displayName !== undefined) {
      update.displayName = data.displayName;
    }
    if (data.newPassword) {
      if (!data.currentPassword) {
        throw new BadRequestException('Current password required to set new password');
      }
      const valid = await bcrypt.compare(data.currentPassword, user.passwordHash);
      if (!valid) throw new BadRequestException('Current password is incorrect');
      update.passwordHash = await bcrypt.hash(data.newPassword, 10);
    }

    if (Object.keys(update).length === 0) {
      return { message: 'No changes' };
    }

    await this.prisma.user.update({ where: { id: userId }, data: update });
    return { message: 'Profile updated' };
  }

  async totpSetup(userId: number) {
    const user = await this.prisma.user.findUnique({ where: { id: userId } });
    if (!user) throw new UnauthorizedException();
    if (user.totpEnabled) throw new BadRequestException('TOTP already enabled');

    const secret = speakeasy.generateSecret({
      name: `Wraith (${user.email})`,
      issuer: 'Wraith',
    });

    // Store secret temporarily (not enabled until verified)
    await this.prisma.user.update({
      where: { id: userId },
      data: { totpSecret: secret.base32 },
    });

    const qrCodeUrl = await QRCode.toDataURL(secret.otpauth_url!);

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

  async totpVerify(userId: number, code: string) {
    const user = await this.prisma.user.findUnique({ where: { id: userId } });
    if (!user || !user.totpSecret) throw new BadRequestException('TOTP not set up');

    const verified = speakeasy.totp.verify({
      secret: user.totpSecret,
      encoding: 'base32',
      token: code,
      window: 1,
    });

    if (!verified) throw new BadRequestException('Invalid code — try again');

    await this.prisma.user.update({
      where: { id: userId },
      data: { totpEnabled: true },
    });

    return { message: 'TOTP enabled successfully' };
  }

  async totpDisable(userId: number, password: string) {
    const user = await this.prisma.user.findUnique({ where: { id: userId } });
    if (!user) throw new UnauthorizedException();

    const valid = await bcrypt.compare(password, user.passwordHash);
    if (!valid) throw new BadRequestException('Password incorrect');

    await this.prisma.user.update({
      where: { id: userId },
      data: { totpEnabled: false, totpSecret: null },
    });

    return { message: 'TOTP disabled' };
  }

  // ─── Admin User Management ──────────────────────────────────────────────

  async listUsers() {
    return this.prisma.user.findMany({
      select: { id: true, email: true, displayName: true, role: true, totpEnabled: true, createdAt: true },
      orderBy: { createdAt: 'asc' },
    });
  }

  async createUser(data: { email: string; password: string; displayName?: string; role?: string }) {
    const existing = await this.prisma.user.findUnique({ where: { email: data.email } });
    if (existing) throw new BadRequestException('Email already in use');

    const hash = await bcrypt.hash(data.password, 10);
    const user = await this.prisma.user.create({
      data: {
        email: data.email,
        passwordHash: hash,
        displayName: data.displayName || null,
        role: data.role || 'user',
      },
    });
    return { id: user.id, email: user.email, displayName: user.displayName, role: user.role };
  }

  async adminUpdateUser(targetId: number, adminId: number, data: { email?: string; displayName?: string; role?: string }) {
    const target = await this.prisma.user.findUnique({ where: { id: targetId } });
    if (!target) throw new NotFoundException('User not found');

    // Prevent removing your own admin role
    if (targetId === adminId && data.role && data.role !== 'admin') {
      throw new ForbiddenException('Cannot remove your own admin role');
    }

    return this.prisma.user.update({
      where: { id: targetId },
      data: {
        email: data.email,
        displayName: data.displayName,
        role: data.role,
      },
      select: { id: true, email: true, displayName: true, role: true },
    });
  }

  async adminResetPassword(targetId: number, newPassword: string) {
    const target = await this.prisma.user.findUnique({ where: { id: targetId } });
    if (!target) throw new NotFoundException('User not found');

    const hash = await bcrypt.hash(newPassword, 10);
    await this.prisma.user.update({ where: { id: targetId }, data: { passwordHash: hash } });
    return { message: 'Password reset' };
  }

  async adminResetTotp(targetId: number) {
    const target = await this.prisma.user.findUnique({ where: { id: targetId } });
    if (!target) throw new NotFoundException('User not found');

    await this.prisma.user.update({
      where: { id: targetId },
      data: { totpEnabled: false, totpSecret: null },
    });
    return { message: 'TOTP reset' };
  }

  async adminDeleteUser(targetId: number, adminId: number) {
    if (targetId === adminId) {
      throw new ForbiddenException('Cannot delete your own account');
    }
    const target = await this.prisma.user.findUnique({ where: { id: targetId } });
    if (!target) throw new NotFoundException('User not found');

    // CASCADE will delete all their hosts, credentials, keys, logs
    await this.prisma.user.delete({ where: { id: targetId } });
    return { message: 'User deleted' };
  }
}