import { SRP, SrpClient } from "fast-srp-hap";
import type { ApiClient } from "../api/client.js";
import { deriveKEK, deriveLoginSubkey, fromBase64 } from "../crypto/index.js";
import type {
    AuthorizationResponse,
    LoginChallenge,
    SRPAttributes,
} from "./types.js";

// Ente uses the 4096-bit SRP group. Both client and server must agree.
const SRP_PARAMS = SRP.params["4096"];

// Wrap fast-srp-hap's callback-based genKey into a Promise.
const genKey = (): Promise<Buffer> => SRP.genKey();

// Create an SRP client from the derived login subkey and SRP attributes.
const makeSrpClient = async (
    srpSalt: Buffer,
    srpUserID: string,
    loginSubKey: Uint8Array,
): Promise<SrpClient> => {
    const clientKey = await genKey();
    return new SrpClient(
        SRP_PARAMS,
        srpSalt,
        Buffer.from(srpUserID),
        Buffer.from(loginSubKey),
        clientKey,
        false,
    );
};

// Perform the two-round SRP handshake against the Ente API.
const srpHandshake = async (
    api: ApiClient,
    srpAttributes: SRPAttributes,
    password: string,
): Promise<LoginChallenge> => {
    const kek = await deriveKEK(
        password,
        fromBase64(srpAttributes.kekSalt),
        srpAttributes.opsLimit,
        srpAttributes.memLimit,
    );
    const loginSubKey = deriveLoginSubkey(kek);

    const srpClient = await makeSrpClient(
        Buffer.from(fromBase64(srpAttributes.srpSalt)),
        srpAttributes.srpUserID,
        loginSubKey,
    );

    const srpA = srpClient.computeA().toString("base64");

    const { sessionID, srpB } = await api.postJSON<{
        sessionID: string;
        srpB: string;
    }>("/users/srp/create-session", {
        srpUserID: srpAttributes.srpUserID,
        srpA,
    });

    srpClient.setB(Buffer.from(srpB, "base64"));
    const srpM1 = srpClient.computeM1().toString("base64");

    const verifyResponse = await api.postJSON<
        AuthorizationResponse & { srpM2: string }
    >("/users/srp/verify-session", {
        sessionID,
        srpUserID: srpAttributes.srpUserID,
        srpM1,
    });

    srpClient.checkM2(Buffer.from(verifyResponse.srpM2, "base64"));

    if (verifyResponse.twoFactorSessionID) {
        return { kind: "totp", sessionID: verifyResponse.twoFactorSessionID };
    }
    if (verifyResponse.passkeySessionID) {
        return {
            kind: "passkey",
            sessionID: verifyResponse.passkeySessionID,
        };
    }

    return { kind: "complete", response: verifyResponse };
};

export const beginLogin = async (
    api: ApiClient,
    email: string,
    password: string,
): Promise<LoginChallenge> => {
    const { attributes } = await api.getJSON<{
        attributes: SRPAttributes;
    }>("/users/srp/attributes", { email });

    if (attributes.isEmailMFAEnabled) {
        return { kind: "emailOTP" };
    }

    return srpHandshake(api, attributes, password);
};

export const submitTOTP = async (
    api: ApiClient,
    sessionID: string,
    code: string,
): Promise<AuthorizationResponse> =>
    api.postJSON<AuthorizationResponse>("/users/two-factor/verify", {
        sessionID,
        code,
    });

export const requestEmailOTP = async (
    api: ApiClient,
    email: string,
): Promise<void> => {
    await api.postJSON("/users/ott", { email, purpose: "login" });
};

export const submitEmailOTP = async (
    api: ApiClient,
    email: string,
    code: string,
): Promise<AuthorizationResponse> =>
    api.postJSON<AuthorizationResponse>("/users/verify-email", {
        email,
        ott: code,
    });