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 CohereHandler extends BaseHandler {
	private baseUrl = 'https://api.cohere.com/v2';

	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('cohere');

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

			const messages = this.prepareMessages(options);

			const payload: any = {
				model: config.model,
				messages,
				stream: isStream,
			};

			if (config.temperature !== undefined) payload.temperature = config.temperature;
			if (config.top_p !== undefined) payload.p = config.top_p;
			if (config.max_tokens !== undefined) payload.max_tokens = config.max_tokens;
			if (config.stop) payload.stop_sequences = config.stop;

			if (isStream) {
				return await this.executeStream(payload, apiKey, isPlain);
			} else {
				return await this.executeSync(payload, apiKey, isPlain);
			}
		} catch (error) {
			return this.handleError(error, isPlain);
		}
	}

	getConfig(generalConfig: IPromptGenericConfig): any {
		const rawModel = generalConfig.model || 'command-r-plus';
		const model = this.normalizeModelName(rawModel, 'cohere');
		const maxTokens = generalConfig.maxTokens || 4096;

		return {
			model,
			temperature: generalConfig.temperature ?? 0.7,
			max_tokens: maxTokens,
			top_p: generalConfig.topP,
			stop: generalConfig.stopSequences,
		};
	}

	private prepareMessages(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) {
				const message: any = {
					role: msg.role === 'assistant' ? 'assistant' : msg.role,
					content: msg.content,
				};

				messages.push(message);
			}
		}

		return messages;
	}

	private async executeStream(payload: any, apiKey: string, isPlain: boolean): Promise<ReadableStream> {
		const { readable, writable } = new TransformStream();
		const writer = writable.getWriter();
		const encoder = new TextEncoder();

		const url = `${this.baseUrl}/chat`;

		libx.log.v('CohereHandler: streaming request to', url);

		let faultyChunk = '';
		const promise = Streams.getStream(
			url,
			async (chunk: string) => {
				try {
					// Cohere uses SSE format
					let incoming = faultyChunk + chunk;
					faultyChunk = '';

					const events = incoming.split(/\n\n+/);
					if (!incoming.endsWith('\n\n')) {
						faultyChunk = events.pop() || '';
					}

					for (const evt of events) {
						const lines = evt.split('\n').map(l => l.trim()).filter(Boolean);
						if (lines.length === 0) continue;

						for (let line of lines) {
							if (!line.startsWith('data:')) continue;
							const jsonStr = line.slice(5).trim();
							if (!jsonStr) continue;

							try {
								const data = JSON.parse(jsonStr);
								// Cohere v2 API streaming format
								if (data.type === 'content-delta') {
									const delta = data.delta?.message?.content?.text;
									if (delta) {
										await this.writeStreamChunk(writer, encoder, delta, isPlain);
									}
								}
							} catch (e) {
								faultyChunk += line + '\n';
							}
						}
					}
				} catch (err) {
					libx.log.w('CohereHandler: chunk processing error', err);
				}
			},
			{
				method: 'POST',
				headers: {
					'Content-Type': 'application/json',
					'Authorization': `Bearer ${apiKey}`,
					'Accept': 'text/event-stream'
				},
				body: payload,
				encoding: 'utf-8',
				useEventBuffering: false,
			}
		);

		promise.then(() => {
			libx.log.v('CohereHandler: stream completed');
			if (!isPlain) {
				writer.write(encoder.encode('data: [DONE]\n\n'));
			}
			writer.close();
		}).catch((error) => {
			libx.log.e('CohereHandler: stream error', error);
			if (isPlain) {
				writer.write(encoder.encode(`\nError: ${error.message}`));
			} else {
				writer.write(encoder.encode(`data: ${JSON.stringify({ error: error.message })}\n\n`));
			}
			writer.close();
		});

		return readable;
	}

	private async executeSync(payload: any, apiKey: string, isPlain: boolean): Promise<ReadableStream> {
		const url = `${this.baseUrl}/chat`;
		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(`Cohere API error: ${error}`);
		}
		const data = await response.json();
		const content = data.message?.content?.[0]?.text || '';

		const metadata = {
			model: data.model,
			usage: data.usage,
			finish_reason: data.finish_reason
		};

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