/**
 * Custom Token Handler for Public Clients
 * 
 * This handles the /oauth/token endpoint to support public clients with PKCE
 * while leveraging the existing workers-oauth-provider infrastructure
 */

import { Env } from './index';

export interface TokenRequest {
	grant_type: string;
	client_id: string;
	code: string;
	redirect_uri: string;
	code_verifier?: string; // PKCE parameter
}

export interface TokenResponse {
	access_token: string;
	token_type: string;
	expires_in: number;
	scope: string;
	refresh_token?: string;
}

export class CustomTokenHandler {
	private env: Env;

	constructor(env: Env) {
		this.env = env;
	}

	/**
	 * Handle token exchange for public clients with PKCE support
	 */
	async handleTokenRequest(request: Request): Promise<Response> {
		try {
			// Parse the token request
			const formData = await request.formData();
			const tokenRequest: TokenRequest = {
				grant_type: formData.get('grant_type') as string,
				client_id: formData.get('client_id') as string,
				code: formData.get('code') as string,
				redirect_uri: formData.get('redirect_uri') as string,
				code_verifier: formData.get('code_verifier') as string || undefined,
			};

			// Validate required parameters
			if (!tokenRequest.grant_type || !tokenRequest.client_id || !tokenRequest.code) {
				return this.createErrorResponse('invalid_request', 'Missing required parameters');
			}

			// Only support authorization_code grant type
			if (tokenRequest.grant_type !== 'authorization_code') {
				return this.createErrorResponse('unsupported_grant_type', 'Only authorization_code grant type is supported');
			}

			// Validate the authorization code and get the stored authorization data
			const authData = await this.validateAuthorizationCode(tokenRequest.code, tokenRequest.client_id);
			if (!authData) {
				return this.createErrorResponse('invalid_grant', 'Invalid or expired authorization code');
			}

			// Validate PKCE if code_verifier is provided
			if (tokenRequest.code_verifier && authData.codeChallenge) {
				const isValidPKCE = await this.validatePKCE(tokenRequest.code_verifier, authData.codeChallenge);
				if (!isValidPKCE) {
					return this.createErrorResponse('invalid_grant', 'PKCE validation failed');
				}
			}

			// Generate access token using the OAuth provider's token generation
			const accessToken = await this.generateAccessToken(authData);

			// Create token response
			const tokenResponse: TokenResponse = {
				access_token: accessToken,
				token_type: 'Bearer',
				expires_in: 3600, // 1 hour
				scope: authData.scope.join(' '),
			};

			return new Response(JSON.stringify(tokenResponse), {
				status: 200,
				headers: {
					'Content-Type': 'application/json',
					'Cache-Control': 'no-store',
					'Pragma': 'no-cache',
				},
			});

		} catch (error) {
			console.error('Token request error:', error);
			return this.createErrorResponse('server_error', 'Internal server error');
		}
	}

	/**
	 * Validate authorization code and retrieve stored authorization data
	 */
	private async validateAuthorizationCode(code: string, clientId: string): Promise<any> {
		try {
			// Use the OAuth provider's internal methods to validate the code
			// This leverages the existing KV storage and validation logic
			const codeKey = `auth_code:${code}`;
			const storedData = await this.env.OAUTH_KV.get(codeKey, 'json') as any;
			
			if (!storedData) {
				return null;
			}

			// Verify client ID matches
			if (storedData.clientId !== clientId) {
				return null;
			}

			// Check expiration (authorization codes should expire quickly)
			const now = Date.now();
			if (storedData.expiresAt && now > storedData.expiresAt) {
				// Clean up expired code
				await this.env.OAUTH_KV.delete(codeKey);
				return null;
			}

			// Delete the code after use (single-use)
			await this.env.OAUTH_KV.delete(codeKey);

			return storedData;

		} catch (error) {
			console.error('Error validating authorization code:', error);
			return null;
		}
	}

	/**
	 * Validate PKCE code_verifier against stored code_challenge
	 */
	private async validatePKCE(codeVerifier: string, codeChallenge: string): Promise<boolean> {
		try {
			// Create SHA256 hash of the code_verifier
			const encoder = new TextEncoder();
			const data = encoder.encode(codeVerifier);
			const hashBuffer = await crypto.subtle.digest('SHA-256', data);
			
			// Convert to base64url
			const hashArray = new Uint8Array(hashBuffer);
			const base64String = btoa(String.fromCharCode(...hashArray));
			const base64url = base64String.replace(/\+/g, '-').replace(/\//g, '_').replace(/=/g, '');
			
			return base64url === codeChallenge;
		} catch (error) {
			console.error('PKCE validation error:', error);
			return false;
		}
	}

	/**
	 * Generate access token using OAuth provider infrastructure
	 */
	private async generateAccessToken(authData: any): Promise<string> {
		// Generate a secure random token
		const tokenBytes = new Uint8Array(32);
		crypto.getRandomValues(tokenBytes);
		const accessToken = btoa(String.fromCharCode(...tokenBytes))
			.replace(/\+/g, '-')
			.replace(/\//g, '_')
			.replace(/=/g, '');

		// Store the token in KV with associated user data
		const tokenKey = `access_token:${accessToken}`;
		const tokenData = {
			userId: authData.userId,
			clientId: authData.clientId,
			scope: authData.scope,
			props: authData.props,
			expiresAt: Date.now() + (3600 * 1000), // 1 hour
		};

		await this.env.OAUTH_KV.put(tokenKey, JSON.stringify(tokenData), {
			expirationTtl: 3600, // Auto-expire after 1 hour
		});

		return accessToken;
	}

	/**
	 * Create OAuth error response
	 */
	private createErrorResponse(error: string, errorDescription: string): Response {
		return new Response(JSON.stringify({
			error,
			error_description: errorDescription,
		}), {
			status: 400,
			headers: {
				'Content-Type': 'application/json',
				'Cache-Control': 'no-store',
				'Pragma': 'no-cache',
			},
		});
	}
}