Personal-Test-Ground/src/auth/auth.controller.ts

import { BadRequestException, Body, ConflictException, Controller, Get, Logger, NotFoundException, Post, Session, UseGuards } from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import { AuthenticatorTransportFuture } from '@simplewebauthn/server';
import { isoBase64URL } from '@simplewebauthn/server/helpers';
import { JwtAuthGuard } from 'src/common/guards/jwt-auth.guard';
import { serverConfig } from 'src/config';
import { UsersService } from 'src/users/user.service';
import { WebauthnService } from 'src/auth/webauthn.service';
import { LoginPayload, User } from 'src/interface/interface';

@Controller('auth')
export class AuthController {
    private logger: Logger = new Logger(`Auth Controller`)
    constructor(private webauthnService: WebauthnService, private usersService: UsersService, private jwtService: JwtService) { }

    @Post('webauthn-register-options')
    async getRegistrationOptions(
        @Body() body: { username: string },
        @Session() session: Record<string, any>,
    ) {
        this.logger.log(`Registering options...`)
        if (!body.username?.trim()) {
            throw new BadRequestException('Missing or empty username');
        }

        const user = await this.getUserFromDb(body.username);
        if (!user) {
            throw new NotFoundException('User not found');
        }

        const options = await this.webauthnService.generateRegistrationOptions(user);
        session.challenge = options.challenge;
        return options;
    }


    @Post('webauthn-register')
    async registerCredential(
        @Body() body: any,
        @Session() session: Record<string, any>,
    ) {
        this.logger.log(`Registering ${body.name}`)
        const expectedChallenge = session.challenge;
        if (!expectedChallenge) {
            throw new BadRequestException('Missing challenge in session');
        }

        // 1. Verify the registration response
        const verification = await this.webauthnService.verifyRegistrationResponse(body, expectedChallenge);
        const { verified, registrationInfo } = verification;

        if (!verified || !registrationInfo) {
            throw new BadRequestException('Registration verification failed');
        }

        // 2. Check if user exists
        let user = await this.usersService.findByEmail(body.email);

        if (!user) {
            // 👇 Create new user if not found
            user = await this.usersService.createUser({
                name: body.name,
                email: body.email,
                password: '', // or null/undefined if you're not using passwords for passkey
            });
        } else {
            // 👇 Optional: prevent duplicate registration
            if (user.devices?.length) {
                throw new ConflictException('User already registered with WebAuthn');
            }
        }

        // 3. Sanitize transports
        const rawTransports = body.response?.transports ?? [];
        const allowedTransports: AuthenticatorTransport[] = ['usb', 'ble', 'nfc', 'internal'];

        const transports = (rawTransports as AuthenticatorTransportFuture[]).filter(
            (t): t is AuthenticatorTransport => allowedTransports.includes(t as AuthenticatorTransport),
        );

        // Optional: Log dropped transports
        const dropped = rawTransports.filter(
            t => !allowedTransports.includes(t as AuthenticatorTransport),
        );
        if (dropped.length) {
            console.warn('Dropped unsupported transports:', dropped);
        }

        // 4. Store device
        user.devices = user.devices || [];
        user.devices.push({
            credentialID: isoBase64URL.toBuffer(registrationInfo.credential.id),
            credentialPublicKey: registrationInfo.credential.publicKey,
            counter: registrationInfo.credential.counter,
            transports, // ✅ now correct
        });

        // 5. Create JWT and return
        const payload = { sub: user.id, name: user.name, email: user.email };
        const token = this.jwtService.sign(payload);
        let response = {
            verified: true,
            access_token: token,
            name: body.name,
        };

        return response
    }

    @Post('webauthn-login-options')
    async getLoginOptions(
        @Body() body: { email?: string; passkey?: boolean },
        @Session() session: Record<string, any>,
    ) {
        if (body.passkey) {
            this.logger.log(`Generating passkey login options...`);

            const options = await this.webauthnService.generateAuthenticationOptions(); // No allowCredentials
            session.challenge = options.challenge;
            return options;
        }

        if (!body.email) {
            throw new BadRequestException('Missing email');
        }

        this.logger.log(`Generating login options for ${body.email}...`);

        const user = await this.usersService.findByEmail(body.email);
        if (!user || !user.devices?.length) {
            throw new NotFoundException('User or devices not found');
        }

        const options = await this.webauthnService.generateAuthenticationOptions(user.devices);
        session.challenge = options.challenge;
        session.email = user.email;
        return options;
    }

    @Post('webauthn-login')
    async verifyLogin(@Body() body: any, @Session() session: Record<string, any>) {
        const expectedChallenge = session.challenge;
        if (!expectedChallenge) {
            throw new BadRequestException('Missing challenge in session');
        }

        let user;

        // Case 1: email-based login
        if (session.email) {
            user = await this.usersService.findByEmail(session.email);
        } else {
            // Case 2: passkey login (discoverable credential)
            const credentialId = body.id || body.rawId;
            user = await this.usersService.findByCredentialId(credentialId);
        }

        if (!user || !user.devices?.length) {
            throw new NotFoundException('User or devices not found');
        }

        const credentialId = body.id || body.rawId;
        const device = user.devices.find(
            d => isoBase64URL.fromBuffer(d.credentialID) === credentialId,
        );

        if (!device) {
            throw new NotFoundException('Device not registered');
        }

        const verification = await this.webauthnService.verifyAuthenticationResponse(
            body,
            expectedChallenge,
            device.credentialPublicKey,
            isoBase64URL.fromBuffer(device.credentialID),
            device.counter,
        );

        if (!verification.verified) {
            throw new BadRequestException('Authentication failed');
        }

        device.counter = verification.authenticationInfo.newCounter;
        await this.usersService.updateUser(user);

        const payload = { sub: user.id, name: user.name, email: user.email };
        const token = this.jwtService.sign(payload);

        return {
            verified: true,
            access_token: token,
            name: user.name,
        };
    }


    @Post('auth-options')
    async getAuthenticationOptions(@Body() body: any, @Session() session: Record<string, any>) {
        const devices = await this.getDevicesFromDb(body.userID);
        const options = await this.webauthnService.generateAuthenticationOptions(devices);
        session.challenge = options.challenge;
        return options;
    }

    @Post('webauthn-authenticate')
    async verifyAuthentication(@Body() body: any, @Session() session: Record<string, any>) {
        const user = await this.getUserFromDb(body.userID);
        const device = await this.getDeviceByCredentialID(body.rawId);

        const verification = await this.webauthnService.verifyAuthenticationResponse(
            body,
            session.challenge,
            device.credentialPublicKey,
            device.credentialID,
            device.counter,
        );

        if (verification.verified) {
            // Update counter in DB
        }

        return verification;
    }

    @Post('register')
    async register(@Body() body: User) {
        this.logger.log(`[REGISTER] Incoming registration request for: ${body.email}`);

        if (!body.email || !body.password || !body.name) {
            this.logger.warn(`[REGISTER] Missing required fields`);
            throw new BadRequestException('Missing required fields');
        }

        let user;
        try {
            user = await this.usersService.createUser(body);
            this.logger.log(`[REGISTER] ✅ User created successfully: ${user.email}`);
        } catch (err) {
            this.logger.error(`[REGISTER] ❌ Failed to create user: ${err.message}`);
            throw new ConflictException(err);
        }

        // Generate JWT after successful registration
        const payload = { sub: user.id, name: user.name, email: user.email };
        const token = this.jwtService.sign(payload);

        this.logger.debug(`[REGISTER] Returning success response for: ${body.email}`);
        return {
            verified: true,
            access_token: token,
            name: user.name,
        };
    }

    @Post('login')
    async login(@Body() body: LoginPayload) {
        this.logger.log(`[LOGIN] Incoming login request for: ${body.email}`);

        if (!body.email || !body.password) {
            this.logger.warn(`[LOGIN] Missing email or password`);
            throw new BadRequestException('Missing email or password');
        }

        let user;
        try {
            user = await this.usersService.validateUser(body);
            this.logger.log(`[LOGIN] ✅ User validated: ${user.email}`);
        } catch (err) {
            this.logger.error(`[LOGIN] ❌ Invalid credentials or error: ${err.message}`);
            throw new BadRequestException(err);
        }

        // Generate JWT for login as well
        const payload = { sub: user.id, name: user.name, email: user.email };
        const token = this.jwtService.sign(payload);

        this.logger.debug(`[LOGIN] Returning success response for: ${body.email}`);
        return {
            verified: true,
            access_token: token,
            name: user.name,
        };
    }


    @Get(`server`)
    @UseGuards(JwtAuthGuard)
    getServerAPIurl(): string {
        this.logger.log(`attempted acquisition`)
        return serverConfig.exposedUrl
    }


    // -- Mock DB lookups (replace with your DB calls) --
    async getUserFromDb(userID: string) {
        return {
            id: userID,
            username: `${userID}@user.com`,
            devices: [], // Array of previously registered devices
        };
    }

    async getDevicesFromDb(userID: string) {
        return []; // Retrieve from DB
    }

    async getDeviceByCredentialID(id: string) {
        return {
            credentialID: id,
            credentialPublicKey: new Uint8Array(), // from DB
            counter: 0,
        };
    }
}