import {
    ALGORITHMS,
    encodePublicKeychain,
    extractPublicKeys,
    type GenerateKeychainOptions,
    generateKeychain,
    type Keychain,
    type KeySlotPrivate,
} from "@suaveplan/crypto/keychain";
import type { Clock } from "@suaveplan/types";

// Mirror of `LayoutV1Fields` from `@suaveplan/crypto/derivation/layouts`.
// Re-declared locally because the type is not re-exported from
// `@suaveplan/crypto/keychain`; the structural shape matches and
// `generateKeychain` accepts any object satisfying it.
interface LayoutV1Fields {
    readonly masterVersion: number;
    readonly userId: number;
    readonly depth: number;
    readonly slotVersion: number;
}

import {
    CHAIN_KEYS_CODES,
    ChainKeysError,
} from "../error-codes/error-codes.js";
import { ChainedKeySigner } from "../signer/signer.js";

// Branded purpose-tag keys. The genesis crypto `Keychain` is
// `Record<PurposeTag, KeySlotPrivate>` where `PurposeTag` is a branded
// `string & { __brand }` type; cast literals once via this typed
// constant so downstream access typechecks.
const MAST = "mast" as unknown as keyof Keychain<LayoutV1Fields>;
const SIGN = "sign" as unknown as keyof Keychain<LayoutV1Fields>;

function getSlot(
    keys: Keychain<LayoutV1Fields>,
    tag: keyof Keychain<LayoutV1Fields>
): KeySlotPrivate<LayoutV1Fields> | undefined {
    return (keys as Record<string, KeySlotPrivate<LayoutV1Fields>>)[
        tag as unknown as string
    ];
}

const DEFAULT_LAYOUT_FIELDS: LayoutV1Fields = {
    masterVersion: 1,
    userId: 0,
    depth: 0,
    slotVersion: 1,
};

const DEFAULT_CLOCK: Clock = () => Date.now();

/**
 * Result of every factory call: the generated `Keychain` plus the signed
 * chain bytes that bind its public projection to the parent (or to itself
 * for a root).
 */
export interface FactoryResult {
    readonly keys: Keychain<LayoutV1Fields>;
    readonly chain: Uint8Array;
}

/** Options for `createRootKeysWithChain`. */
export interface CreateRootOptions {
    /** Inject a pre-configured signer; defaults to `new ChainedKeySigner()`. */
    readonly signer?: ChainedKeySigner;
    /** Override the keychain generation options. */
    readonly generate?: Partial<GenerateKeychainOptions>;
}

/** Options for `createServerKeysWithChain`. */
export interface CreateServerOptions {
    readonly rootKeys: Keychain<LayoutV1Fields>;
    readonly rootChain: Uint8Array;
    readonly signer?: ChainedKeySigner;
    readonly generate?: Partial<GenerateKeychainOptions>;
}

/** Options for `rollServerKeysWithChain`. */
export interface RollServerOptions {
    readonly rootKeys: Keychain<LayoutV1Fields>;
    readonly rootChain: Uint8Array;
    readonly currentVersion: number;
    readonly signer?: ChainedKeySigner;
}

/**
 * Generate a fresh root keychain and self-sign its public projection.
 */
export async function createRootKeysWithChain(
    options: CreateRootOptions = {}
): Promise<FactoryResult> {
    const signer = options.signer ?? new ChainedKeySigner();
    const keys = generateKeychain(buildGenerateOptions(options.generate));
    const publicKeys = extractPublicKeys(keys);
    const payload = encodePublicKeychain(publicKeys);
    const signSlot = getSlot(keys, SIGN) as KeySlotPrivate<LayoutV1Fields>;
    const chain = await signer.signWithSlot(payload, signSlot);
    return { keys, chain };
}

/**
 * Derive a server keychain from `rootKeys.master`, sign its public projection
 * with `rootKeys.sign`, and attach the parent `rootChain`.
 */
export async function createServerKeysWithChain(
    options: CreateServerOptions
): Promise<FactoryResult> {
    const signer = options.signer ?? new ChainedKeySigner();
    assertRootSlots(options.rootKeys);
    const generateOptions = buildGenerateOptions({
        ...options.generate,
        masterKey: (
            getSlot(options.rootKeys, MAST) as KeySlotPrivate<LayoutV1Fields>
        ).secretKey,
    });
    const keys = generateKeychain(generateOptions);
    const publicKeys = extractPublicKeys(keys);
    const payload = encodePublicKeychain(publicKeys);
    const chain = await signer.signWithSlot(
        payload,
        getSlot(options.rootKeys, SIGN) as KeySlotPrivate<LayoutV1Fields>,
        options.rootChain
    );
    return { keys, chain };
}

/**
 * Roll the server keychain forward to the next version, re-signed by
 * `rootKeys.sign`. The new keychain shares the master key with the
 * existing root so the derivation chain stays intact.
 */
export async function rollServerKeysWithChain(
    options: RollServerOptions
): Promise<FactoryResult> {
    if (
        !Number.isInteger(options.currentVersion) ||
        options.currentVersion < 0
    ) {
        throw new ChainKeysError(
            CHAIN_KEYS_CODES.INVALID_KEYCHAIN,
            `currentVersion must be a non-negative integer (got ${options.currentVersion})`
        );
    }
    const signer = options.signer ?? new ChainedKeySigner();
    assertRootSlots(options.rootKeys);
    const generateOptions = buildGenerateOptions({
        masterKey: (
            getSlot(options.rootKeys, MAST) as KeySlotPrivate<LayoutV1Fields>
        ).secretKey,
        layoutFields: {
            ...DEFAULT_LAYOUT_FIELDS,
            slotVersion: options.currentVersion + 1,
        },
    });
    const keys = generateKeychain(generateOptions);
    const publicKeys = extractPublicKeys(keys);
    const payload = encodePublicKeychain(publicKeys);
    const chain = await signer.signWithSlot(
        payload,
        getSlot(options.rootKeys, SIGN) as KeySlotPrivate<LayoutV1Fields>,
        options.rootChain
    );
    return { keys, chain };
}

function buildGenerateOptions(
    overrides: Partial<GenerateKeychainOptions> | undefined
): GenerateKeychainOptions {
    const base: GenerateKeychainOptions = {
        layoutFields: DEFAULT_LAYOUT_FIELDS,
        clock: DEFAULT_CLOCK,
    };
    if (!overrides) return base;
    return {
        ...base,
        ...overrides,
        layoutFields: overrides.layoutFields ?? DEFAULT_LAYOUT_FIELDS,
        clock: overrides.clock ?? DEFAULT_CLOCK,
    };
}

function assertRootSlots(rootKeys: Keychain<LayoutV1Fields>): void {
    const master = getSlot(rootKeys, MAST);
    if (
        master === undefined ||
        !(master.secretKey instanceof Uint8Array) ||
        master.secretKey.length === 0
    ) {
        throw new ChainKeysError(
            CHAIN_KEYS_CODES.INVALID_KEYCHAIN,
            "rootKeys.master must contain a non-empty secretKey"
        );
    }
    const sign = getSlot(rootKeys, SIGN);
    if (
        sign === undefined ||
        sign.algorithm !== ALGORITHMS.ED25519 ||
        !(sign.secretKey instanceof Uint8Array) ||
        sign.secretKey.length === 0
    ) {
        throw new ChainKeysError(
            CHAIN_KEYS_CODES.INVALID_KEYCHAIN,
            "rootKeys.sign must be a valid Ed25519 slot with non-empty secretKey"
        );
    }
}