import z from "zod";
import { SRP_CODES, srpError } from "../errors/errors.js";

const uint8 = z.custom<Uint8Array>((v) => v instanceof Uint8Array, {
    message: "expected Uint8Array",
});

/**
 * Zod boundary schema for `SrpParams`. Validates byte-shape and
 * algorithm tag; mismatches throw `SrpError(SRP.INVALID_PARAMS)`
 * via `parseSrpParams`.
 */
export const SrpParamsSchema = z
    .object({
        N: uint8,
        g: uint8,
        bitLength: z.number().int().positive(),
        hash: z.object({
            algorithm: z.string(),
            digestSize: z.number().int().positive(),
            digest: z.custom<
                (
                    input: Uint8Array,
                    signal?: AbortSignal
                ) => Uint8Array | Promise<Uint8Array>
            >((v) => typeof v === "function", {
                message: "hash.digest must be a function",
            }),
        }),
        algorithm: z.union([z.literal("SHA256"), z.literal("SHA1")]),
    })
    .superRefine((value, ctx) => {
        if (value.N.byteLength % 4 !== 0) {
            ctx.addIssue({
                code: z.ZodIssueCode.custom,
                message: `N.byteLength must be a multiple of 4, got ${value.N.byteLength}`,
            });
        }
        if (value.bitLength !== value.N.byteLength * 8) {
            ctx.addIssue({
                code: z.ZodIssueCode.custom,
                message: "bitLength must equal N.byteLength * 8",
            });
        }
        const gInt = decodeBigEndian(value.g);
        const NInt = decodeBigEndian(value.N);
        if (gInt <= 1n) {
            ctx.addIssue({
                code: z.ZodIssueCode.custom,
                message: `g must decode to a big-int > 1, got ${gInt}`,
            });
        }
        if (gInt >= NInt) {
            ctx.addIssue({
                code: z.ZodIssueCode.custom,
                message: "g must decode to a big-int strictly less than N",
            });
        }
    });

/**
 * Per-handshake state envelope — the subset of values both client and
 * server need to keep in memory across the two round-trips. Validated
 * to catch transport-layer mangling early.
 */
export const SrpHandshakeStateSchema = z.object({
    I: z.string().min(1),
    salt: uint8,
    A: uint8,
    B: uint8,
    role: z.union([z.literal("client"), z.literal("server")]),
});

/** Server-side verifier record: `{ I, salt, v }`. */
export const VerifierSchema = z.object({
    I: z.string().min(1),
    salt: uint8,
    v: uint8,
});

/** Parse + validate an `SrpParams`. Mismatch → `SRP.INVALID_PARAMS`. */
export function parseSrpParams(
    value: unknown
): z.infer<typeof SrpParamsSchema> {
    const result = SrpParamsSchema.safeParse(value);
    if (!result.success) {
        throw srpError(SRP_CODES.INVALID_PARAMS, result.error.message);
    }
    return result.data;
}

function decodeBigEndian(bytes: Uint8Array): bigint {
    let acc = 0n;
    for (const b of bytes) {
        acc = (acc << 8n) | BigInt(b);
    }
    return acc;
}