import { errAsync, okAsync, ResultAsync } from "neverthrow";
import { FlowExecCtx } from "@core/flow.execution.context";
import { hashString, verifyHash } from "@/core/hash.utils";
import { twoFactor, twofaSessions } from "@pkg/db/schema";
import { TwoFactor, type TwoFaSession } from "./data";
import { and, Database, eq, gt, lt } from "@pkg/db";
import { generateSecret, verify } from "otplib";
import { settings } from "@core/settings";
import type { Err } from "@pkg/result";
import { twofaErrors } from "./errors";
import { logger } from "@pkg/logger";
import { Redis } from "@pkg/redis";
import { nanoid } from "nanoid";

type TwoFaSetup = {
    secret: string;
    lastUsedCode: string;
    tries: number;
};

export class TwofaRepository {
    private PENDING_KEY_PREFIX = "pending_enabling_2fa:";
    private EXPIRY_TIME = 60 * 20; // 20 mins
    private DEFAULT_BACKUP_CODES_AMT = 8;
    private MAX_SETUP_ATTEMPTS = 3;

    constructor(
        private db: Database,
        private store: Redis,
    ) {}

    checkTotp(secret: string, code: string) {
        const checked = verify({ secret, token: code });
        logger.debug("TOTP check result", { checked });
        return checked;
    }

    async checkBackupCode(hash: string, code: string) {
        return verifyHash({ hash, target: code });
    }

    private getKey(userId: string) {
        if (userId.includes(this.PENDING_KEY_PREFIX)) {
            return userId;
        }
        return `${this.PENDING_KEY_PREFIX}${userId}`;
    }

    getUsers2FAInfo(
        fctx: FlowExecCtx,
        userId: string,
        returnUndefined?: boolean,
    ): ResultAsync<TwoFactor | undefined, Err> {
        logger.info("Getting user 2FA info", { ...fctx, userId });

        return ResultAsync.fromPromise(
            this.db.query.twoFactor.findFirst({
                where: eq(twoFactor.userId, userId),
            }),
            () => twofaErrors.dbError(fctx, "Failed to query 2FA info"),
        ).andThen((found) => {
            if (!found) {
                logger.debug("2FA info not found for user", {
                    ...fctx,
                    userId,
                });
                if (returnUndefined) {
                    return okAsync(undefined);
                }
                return errAsync(twofaErrors.notFound(fctx));
            }
            logger.info("2FA info retrieved successfully", { ...fctx, userId });
            return okAsync(found as TwoFactor);
        });
    }

    isSetupPending(
        fctx: FlowExecCtx,
        userId: string,
    ): ResultAsync<boolean, Err> {
        logger.debug("Checking if 2FA setup is pending", { ...fctx, userId });

        return ResultAsync.fromPromise(
            this.store.get(this.getKey(userId)),
            () =>
                twofaErrors.dbError(
                    fctx,
                    "Failed to check setup pending status",
                ),
        ).map((found) => {
            const isPending = !!found;
            logger.debug("Setup pending status checked", {
                ...fctx,
                userId,
                isPending,
            });
            return isPending;
        });
    }

    setup(fctx: FlowExecCtx, userId: string): ResultAsync<string, Err> {
        logger.info("Starting 2FA setup", { ...fctx, userId });

        return ResultAsync.fromSafePromise(
            (async () => {
                const secret = generateSecret();
                const payload = {
                    secret,
                    lastUsedCode: "",
                    tries: 0,
                } as TwoFaSetup;
                await this.store.setex(
                    this.getKey(userId),
                    this.EXPIRY_TIME,
                    JSON.stringify(payload),
                );
                logger.info("Created temp 2FA session", {
                    ...fctx,
                    userId,
                    expiresIn: this.EXPIRY_TIME,
                });
                return secret;
            })(),
        ).mapErr(() =>
            twofaErrors.dbError(fctx, "Setting to data store failed"),
        );
    }

    verifyAndEnable2FA(
        fctx: FlowExecCtx,
        userId: string,
        code: string,
    ): ResultAsync<boolean, Err> {
        logger.info("Verifying and enabling 2FA", { ...fctx, userId });

        return ResultAsync.fromPromise(
            this.store.get(this.getKey(userId)),
            () => twofaErrors.dbError(fctx, "Failed to get setup session"),
        )
            .andThen((payload) => {
                if (!payload) {
                    logger.error("Setup session not found", {
                        ...fctx,
                        userId,
                    });
                    return errAsync(twofaErrors.setupNotFound(fctx));
                }
                return okAsync(JSON.parse(payload) as TwoFaSetup);
            })
            .andThen((payloadObj) => {
                const key = this.getKey(userId);

                if (payloadObj.tries >= this.MAX_SETUP_ATTEMPTS) {
                    logger.warn("Max setup attempts reached", {
                        ...fctx,
                        userId,
                        tries: payloadObj.tries,
                    });
                    return ResultAsync.fromPromise(this.store.del(key), () =>
                        twofaErrors.dbError(
                            fctx,
                            "Failed to delete setup session",
                        ),
                    ).andThen(() =>
                        errAsync(twofaErrors.maxAttemptsReached(fctx)),
                    );
                }

                if (
                    !this.checkTotp(payloadObj.secret, code) ||
                    code === payloadObj.lastUsedCode
                ) {
                    logger.warn("Invalid 2FA code during setup", {
                        ...fctx,
                        userId,
                        tries: payloadObj.tries + 1,
                        codeReused: code === payloadObj.lastUsedCode,
                    });
                    return ResultAsync.fromPromise(
                        this.store.setex(
                            key,
                            this.EXPIRY_TIME,
                            JSON.stringify({
                                secret: payloadObj.secret,
                                lastUsedCode: code,
                                tries: payloadObj.tries + 1,
                            }),
                        ),
                        () =>
                            twofaErrors.dbError(
                                fctx,
                                "Failed to update setup session",
                            ),
                    ).map(() => false);
                }

                logger.info("2FA code verified successfully, enabling 2FA", {
                    ...fctx,
                    userId,
                });

                return ResultAsync.fromPromise(this.store.del(key), () =>
                    twofaErrors.dbError(fctx, "Failed to delete setup session"),
                )
                    .andThen(() =>
                        ResultAsync.fromPromise(
                            this.db
                                .insert(twoFactor)
                                .values({
                                    id: nanoid(),
                                    secret: payloadObj.secret,
                                    userId: userId,
                                    createdAt: new Date(),
                                    updatedAt: new Date(),
                                })
                                .execute(),
                            () =>
                                twofaErrors.dbError(
                                    fctx,
                                    "Failed to insert 2FA record",
                                ),
                        ),
                    )
                    .map(() => {
                        logger.info("2FA enabled successfully", {
                            ...fctx,
                            userId,
                        });
                        return true;
                    });
            });
    }

    disable(fctx: FlowExecCtx, userId: string): ResultAsync<boolean, Err> {
        logger.info("Disabling 2FA", { ...fctx, userId });

        return ResultAsync.fromPromise(
            this.db
                .delete(twoFactor)
                .where(eq(twoFactor.userId, userId))
                .execute(),
            () => twofaErrors.dbError(fctx, "Failed to delete 2FA record"),
        ).map((result) => {
            logger.info("2FA disabled successfully", { ...fctx, userId });
            return true;
        });
    }

    generateBackupCodes(
        fctx: FlowExecCtx,
        userId: string,
    ): ResultAsync<string[], Err> {
        logger.info("Generating backup codes", { ...fctx, userId });

        return ResultAsync.fromPromise(
            this.db.query.twoFactor.findFirst({
                where: eq(twoFactor.userId, userId),
            }),
            () => twofaErrors.dbError(fctx, "Failed to query 2FA info"),
        )
            .andThen((found) => {
                if (!found) {
                    logger.error("2FA not enabled for user", {
                        ...fctx,
                        userId,
                    });
                    return errAsync(twofaErrors.backupCodesNotFound(fctx));
                }
                if (found.backupCodes && found.backupCodes.length) {
                    logger.warn("Backup codes already generated", {
                        ...fctx,
                        userId,
                    });
                    return errAsync(
                        twofaErrors.backupCodesAlreadyGenerated(fctx),
                    );
                }
                return okAsync(found);
            })
            .andThen(() => {
                const codes = Array.from(
                    { length: this.DEFAULT_BACKUP_CODES_AMT },
                    () => nanoid(12),
                );

                logger.debug("Backup codes generated, hashing", {
                    ...fctx,
                    userId,
                    count: codes.length,
                });

                return ResultAsync.fromPromise(
                    (async () => {
                        const hashed = [];
                        for (const code of codes) {
                            const hash = await hashString(code);
                            hashed.push(hash);
                        }
                        return { codes, hashed };
                    })(),
                    () =>
                        twofaErrors.dbError(
                            fctx,
                            "Failed to hash backup codes",
                        ),
                ).andThen(({ codes, hashed }) =>
                    ResultAsync.fromPromise(
                        this.db
                            .update(twoFactor)
                            .set({ backupCodes: hashed })
                            .where(eq(twoFactor.userId, userId))
                            .returning(),
                        () =>
                            twofaErrors.dbError(
                                fctx,
                                "Failed to update backup codes",
                            ),
                    ).map(() => {
                        logger.info("Backup codes generated successfully", {
                            ...fctx,
                            userId,
                        });
                        return codes;
                    }),
                );
            });
    }

    get2FASecret(
        fctx: FlowExecCtx,
        userId: string,
    ): ResultAsync<string | null, Err> {
        logger.debug("Getting 2FA secret", { ...fctx, userId });

        return ResultAsync.fromPromise(
            this.db
                .select()
                .from(twoFactor)
                .where(eq(twoFactor.userId, userId))
                .limit(1),
            () => twofaErrors.dbError(fctx, "Failed to query 2FA secret"),
        ).map((result) => {
            if (!result.length) {
                logger.debug("No 2FA secret found", { ...fctx, userId });
                return null;
            }
            logger.debug("2FA secret retrieved", { ...fctx, userId });
            return result[0].secret;
        });
    }

    createSession(
        fctx: FlowExecCtx,
        params: {
            userId: string;
            sessionId: string;
            ipAddress?: string;
            userAgent?: string;
        },
    ): ResultAsync<TwoFaSession, Err> {
        logger.info("Creating 2FA verification session", {
            ...fctx,
            userId: params.userId,
            sessionId: params.sessionId,
        });

        return ResultAsync.fromSafePromise(
            (async () => {
                const expiryMinutes = settings.twofaSessionExpiryMinutes || 10;
                const now = new Date();
                const expiresAt = new Date(
                    now.getTime() + expiryMinutes * 60 * 1000,
                );

                return { expiresAt, now, params };
            })(),
        ).andThen(({ expiresAt, now, params }) =>
            ResultAsync.fromPromise(
                this.db
                    .insert(twofaSessions)
                    .values({
                        id: nanoid(),
                        userId: params.userId,
                        sessionId: params.sessionId,
                        verificationToken: nanoid(32),
                        status: "pending",
                        attempts: 0,
                        maxAttempts: 5,
                        expiresAt,
                        createdAt: now,
                        ipAddress: params.ipAddress,
                        userAgent: params.userAgent,
                    })
                    .returning(),
                () => twofaErrors.dbError(fctx, "Failed to create 2FA session"),
            ).map(([session]) => {
                logger.info("2FA verification session created", {
                    ...fctx,
                    sessionId: session.id,
                    userId: params.userId,
                });
                return session as TwoFaSession;
            }),
        );
    }

    getSessionByToken(
        fctx: FlowExecCtx,
        token: string,
    ): ResultAsync<TwoFaSession | null, Err> {
        logger.debug("Getting 2FA session by token", { ...fctx });

        return ResultAsync.fromPromise(
            this.db
                .select()
                .from(twofaSessions)
                .where(
                    and(
                        eq(twofaSessions.verificationToken, token),
                        gt(twofaSessions.expiresAt, new Date()),
                    ),
                )
                .limit(1),
            () => twofaErrors.dbError(fctx, "Failed to query 2FA session"),
        ).map((result) => {
            if (!result.length) {
                logger.warn("2FA session not found or expired", { ...fctx });
                return null;
            }
            logger.debug("2FA session found", {
                ...fctx,
                sessionId: result[0].id,
            });
            return result[0] as TwoFaSession;
        });
    }

    updateSession(
        fctx: FlowExecCtx,
        id: string,
        updates: Partial<
            Pick<
                TwoFaSession,
                "status" | "attempts" | "verifiedAt" | "codeUsed"
            >
        >,
    ): ResultAsync<TwoFaSession, Err> {
        logger.debug("Updating 2FA session", {
            ...fctx,
            sessionId: id,
            updates,
        });

        return ResultAsync.fromPromise(
            this.db
                .update(twofaSessions)
                .set(updates)
                .where(eq(twofaSessions.id, id))
                .returning(),
            () => twofaErrors.dbError(fctx, "Failed to update 2FA session"),
        ).andThen(([session]) => {
            if (!session) {
                logger.error("2FA session not found for update", {
                    ...fctx,
                    sessionId: id,
                });
                return errAsync(twofaErrors.sessionNotFoundById(fctx));
            }
            logger.debug("2FA session updated successfully", {
                ...fctx,
                sessionId: id,
            });
            return okAsync(session as TwoFaSession);
        });
    }

    incrementAttempts(
        fctx: FlowExecCtx,
        id: string,
    ): ResultAsync<TwoFaSession, Err> {
        logger.debug("Incrementing session attempts", {
            ...fctx,
            sessionId: id,
        });

        return ResultAsync.fromPromise(
            this.db.query.twofaSessions.findFirst({
                where: eq(twofaSessions.id, id),
                columns: { id: true, attempts: true },
            }),
            () =>
                twofaErrors.dbError(
                    fctx,
                    "Failed to query session for increment",
                ),
        )
            .andThen((s) => {
                if (!s) {
                    logger.error("Session not found for increment", {
                        ...fctx,
                        sessionId: id,
                    });
                    return errAsync(twofaErrors.sessionNotFoundById(fctx));
                }
                return okAsync(s);
            })
            .andThen((s) =>
                ResultAsync.fromPromise(
                    this.db
                        .update(twofaSessions)
                        .set({ attempts: s.attempts + 1 })
                        .where(eq(twofaSessions.id, id))
                        .returning(),
                    () =>
                        twofaErrors.dbError(
                            fctx,
                            "Failed to increment attempts",
                        ),
                ).andThen(([session]) => {
                    if (!session) {
                        logger.error("Session not found after increment", {
                            ...fctx,
                            sessionId: id,
                        });
                        return errAsync(twofaErrors.sessionNotFoundById(fctx));
                    }

                    logger.warn("Failed verification attempt", {
                        ...fctx,
                        sessionId: session.id,
                        attempts: session.attempts,
                    });

                    return okAsync(session as TwoFaSession);
                }),
            );
    }

    cleanupExpiredSessions(fctx: FlowExecCtx): ResultAsync<number, Err> {
        logger.info("Cleaning up expired 2FA sessions", { ...fctx });

        return ResultAsync.fromPromise(
            this.db
                .delete(twofaSessions)
                .where(lt(twofaSessions.expiresAt, new Date())),
            () =>
                twofaErrors.dbError(fctx, "Failed to cleanup expired sessions"),
        ).map((result) => {
            const count = result.length || 0;
            logger.info("Expired sessions cleaned up", { ...fctx, count });
            return count;
        });
    }
}