import { BaseHandler } from './BaseHandler';
import { IPromptOptions, IPromptGenericConfig, ModuleOptions } from '../../types';
import { libx } from 'libx.js/build/bundles/essentials.js';
import { Streams } from 'libx.js/build/modules/Streams';

export class AI21Handler extends BaseHandler {
	private baseUrl = 'https://api.ai21.com/studio/v1';

	constructor(options: ModuleOptions) {
		super(options);
	}

	async execute(isStream: boolean, options: IPromptOptions): Promise<ReadableStream> {
		const isPlain = options.config.plain ?? false;
		try {
			const config = this.getConfig(options.config);
			const apiKey = this.getApiKey('ai21');

			if (!apiKey) {
				throw new Error('AI21 API key not configured');
			}

			// Jamba models use chat API, Jurassic-2 models use completion API
			const isJambaModel = config.model.includes('jamba');

			if (isJambaModel) {
				// Use chat API for Jamba models
				return await this.executeChatAPI(isStream, options, config, apiKey, isPlain);
			} else {
				// Use completion API for Jurassic-2 models
				return await this.executeCompletionAPI(options, config, apiKey, isPlain);
			}
		} catch (error) {
			return this.handleError(error, isPlain);
		}
	}

	getConfig(generalConfig: IPromptGenericConfig): any {
		const rawModel = generalConfig.model || 'j2-ultra';
		const model = this.normalizeModelName(rawModel, 'ai21');
		const maxTokens = generalConfig.maxTokens || 2000;

		return {
			model,
			temperature: generalConfig.temperature ?? 0.7,
			max_tokens: maxTokens,
			top_p: generalConfig.topP,
			frequency_penalty: generalConfig.frequencyPenalty || 0,
			presence_penalty: generalConfig.presencePenalty || 0,
			stop: generalConfig.stopSequences,
		};
	}

	private preparePrompt(options: IPromptOptions): string {
		let prompt = '';

		// Add system prompt if provided
		if (options.systemPrompt) {
			prompt += options.systemPrompt + '\n\n';
		}

		// Convert messages to a single prompt string
		if (options.messages) {
			for (const msg of options.messages) {
				if (msg.role === 'system') {
					prompt += msg.content + '\n\n';
				} else if (msg.role === 'user') {
					prompt += `User: ${msg.content}\n`;
				} else if (msg.role === 'assistant') {
					prompt += `Assistant: ${msg.content}\n`;
				}
			}
		}

		// Add final prompt for assistant response
		if (options.messages && options.messages[options.messages.length - 1]?.role === 'user') {
			prompt += 'Assistant:';
		}

		return prompt;
	}

	private async executeChatAPI(isStream: boolean, options: IPromptOptions, config: any, apiKey: string, isPlain: boolean): Promise<ReadableStream> {
		// Jamba models use chat/completions endpoint (OpenAI-compatible)
		const url = `${this.baseUrl}/chat/completions`;

		const messages = this.prepareChatMessages(options);

		const payload: any = {
			model: config.model,
			messages,
			max_tokens: config.max_tokens,
			temperature: config.temperature,
		};

		if (config.top_p !== undefined) payload.top_p = config.top_p;
		if (config.stop) payload.stop = config.stop;

		libx.log.v('AI21Handler (Chat API): request to', url, { model: config.model });

		const response = await fetch(url, {
			method: 'POST',
			headers: {
				'Content-Type': 'application/json',
				'Authorization': `Bearer ${apiKey}`
			},
			body: JSON.stringify(payload)
		});

		if (!response.ok) {
			const error = await response.text();
			throw new Error(`AI21 API error: ${error}`);
		}

		const data = await response.json();
		const content = data.choices?.[0]?.message?.content || '';

		const metadata = {
			model: config.model,
			usage: data.usage,
			finish_reason: data.choices?.[0]?.finish_reason
		};

		return this.writeSyncResponse(content, metadata, isPlain);
	}

	private async executeCompletionAPI(options: IPromptOptions, config: any, apiKey: string, isPlain: boolean): Promise<ReadableStream> {
		// Jurassic-2 models use completion endpoint
		const url = `${this.baseUrl}/${config.model}/complete`;

		const prompt = this.preparePrompt(options);

		const payload: any = {
			prompt,
			maxTokens: config.max_tokens,
			temperature: config.temperature,
		};

		if (config.top_p !== undefined) payload.topP = config.top_p;
		if (config.stop) payload.stopSequences = config.stop;
		if (config.frequency_penalty !== undefined) payload.frequencyPenalty = config.frequency_penalty;
		if (config.presence_penalty !== undefined) payload.presencePenalty = config.presence_penalty;

		libx.log.v('AI21Handler (Completion API): request to', url);

		const response = await fetch(url, {
			method: 'POST',
			headers: {
				'Content-Type': 'application/json',
				'Authorization': `Bearer ${apiKey}`
			},
			body: JSON.stringify(payload)
		});

		if (!response.ok) {
			const error = await response.text();
			throw new Error(`AI21 API error: ${error}`);
		}

		const data = await response.json();
		const content = data.completions?.[0]?.data?.text || '';

		const metadata = {
			model: config.model,
			id: data.id,
			finish_reason: data.completions?.[0]?.finishReason?.reason
		};

		return this.writeSyncResponse(content, metadata, isPlain);
	}

	private prepareChatMessages(options: IPromptOptions): any[] {
		const messages: any[] = [];

		if (options.systemPrompt) {
			messages.push({ role: 'system', content: options.systemPrompt });
		}

		if (options.messages) {
			for (const msg of options.messages) {
				messages.push({
					role: msg.role,
					content: msg.content,
				});
			}
		}

		return messages;
	}
}