import { SignJWT, jwtVerify, exportJWK, generateKeyPair } from 'jose';
import type { OIDCUser } from '../types';

export interface JWTPayload {
    iss: string;
    sub: string;
    aud: string | string[];
    exp: number;
    iat: number;
    [key: string]: any;
}

export interface KeyPair {
    privateKey: any;
    publicKey: any;
    kid: string;
}

/**
 * JWT工具类
 */
export class JWTUtils {
    private privateKey: any;
    private publicKey: any;
    private kid: string;
    private keyPromise: Promise<void> | null = null;

    constructor(signingKey?: string | KeyPair) {
        if (typeof signingKey === 'string') {
            // HMAC密钥
            const encodedKey = new TextEncoder().encode(signingKey);
            this.privateKey = this.publicKey = encodedKey;
            this.kid = 'hmac-key';
        } else if (signingKey) {
            // 非对称密钥
            this.privateKey = signingKey.privateKey;
            this.publicKey = signingKey.publicKey;
            this.kid = signingKey.kid;
        } else {
            // 延迟生成RSA密钥
            this.kid = 'auto-generated';
        }
    }

    /**
     * 确保密钥已初始化
     */
    private async ensureKeys(): Promise<void> {
        if (this.privateKey) return;

        if (!this.keyPromise) {
            this.keyPromise = this.generateRSAKeys();
        }
        await this.keyPromise;
    }

    /**
     * 生成RSA密钥
     */
    private async generateRSAKeys(): Promise<void> {
        const { privateKey, publicKey } = await generateKeyPair('RS256');
        this.privateKey = privateKey;
        this.publicKey = publicKey;
    }

    /**
     * 生成访问令牌
     */
    async generateAccessToken(payload: {
        issuer: string;
        subject: string;
        audience: string;
        clientId: string;
        scope: string;
        expiresIn?: number;
    }): Promise<string> {
        await this.ensureKeys();

        const now = Math.floor(Date.now() / 1000);
        const exp = now + (payload.expiresIn || 3600); // 默认1小时

        return this.signToken({
            iss: payload.issuer,
            sub: payload.subject,
            aud: payload.audience,
            exp,
            iat: now,
            scope: payload.scope,
            client_id: payload.clientId,
            token_type: 'access_token',
        });
    }

    /**
     * 生成ID令牌
     */
    async generateIDToken(payload: {
        issuer: string;
        subject: string;
        audience: string;
        user: OIDCUser;
        authTime?: number;
        nonce?: string;
        expiresIn?: number;
    }): Promise<string> {
        await this.ensureKeys();

        const now = Math.floor(Date.now() / 1000);
        const exp = now + (payload.expiresIn || 3600); // 默认1小时

        const claims = {
            iss: payload.issuer,
            sub: payload.subject,
            aud: payload.audience,
            exp,
            iat: now,
            auth_time: payload.authTime || now,
            ...this.extractUserClaims(payload.user),
        };

        if (payload.nonce) claims.nonce = payload.nonce;

        return this.signToken(claims);
    }

    /**
     * 验证JWT令牌
     */
    async verifyToken(token: string): Promise<JWTPayload> {
        await this.ensureKeys();
        const { payload } = await jwtVerify(token, this.publicKey);
        return payload as JWTPayload;
    }

    /**
     * 生成JWKS
     */
    async generateJWKS(): Promise<{ keys: any[] }> {
        await this.ensureKeys();

        // HMAC密钥不公开
        if (this.kid === 'hmac-key') {
            return { keys: [] };
        }

        const jwk = await exportJWK(this.publicKey);
        return {
            keys: [{
                ...jwk,
                alg: 'RS256',
                use: 'sig',
                kid: this.kid,
            }]
        };
    }

    /**
     * 签名令牌
     */
    private async signToken(payload: any): Promise<string> {
        const algorithm = this.kid === 'hmac-key' ? 'HS256' : 'RS256';
        const header: any = { alg: algorithm };

        if (algorithm === 'RS256') {
            header.kid = this.kid;
        }

        return new SignJWT(payload)
            .setProtectedHeader(header)
            .sign(this.privateKey);
    }

    /**
     * 提取用户声明
     */
    private extractUserClaims(user: OIDCUser): any {
        const claims: any = {};

        if (user.name) claims.name = user.name;
        if (user.given_name) claims.given_name = user.given_name;
        if (user.family_name) claims.family_name = user.family_name;
        if (user.username) claims.preferred_username = user.username;
        if (user.email) claims.email = user.email;
        if (user.email_verified !== undefined) claims.email_verified = user.email_verified;
        if (user.picture) claims.picture = user.picture;
        if (user.phone_number) claims.phone_number = user.phone_number;
        if (user.phone_number_verified !== undefined) claims.phone_number_verified = user.phone_number_verified;

        return claims;
    }

    /**
     * 生成RSA密钥对
     */
    static async generateRSAKeyPair(kid: string = 'rsa-key'): Promise<KeyPair> {
        const { privateKey, publicKey } = await generateKeyPair('RS256');
        return { privateKey, publicKey, kid };
    }
} 