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 GoogleHandler extends BaseHandler {
	private baseUrl = 'https://generativelanguage.googleapis.com/v1beta';

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

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

			// Prepare contents
			const contents = this.prepareContents(options);
			const systemInstruction = options.systemPrompt ? { parts: [{ text: options.systemPrompt }] } : undefined;

			const payload: any = {
				contents,
				generationConfig: {
					temperature: config.temperature,
					maxOutputTokens: config.max_tokens,
					topP: config.top_p,
				},
			};

			if (systemInstruction) {
				payload.systemInstruction = systemInstruction;
			}

			if (config.stop && config.stop.length > 0) {
				payload.generationConfig.stopSequences = config.stop;
			}

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

	getConfig(generalConfig: IPromptGenericConfig): any {
		let model = generalConfig.model || 'gemini-2.0-flash';

		// Strip provider prefix if present
		model = this.normalizeModelName(model, 'google');

		// Normalize model names - if it starts with "models/", keep it; otherwise prepend
		if (!model.startsWith('models/')) {
			model = `models/${model}`;
		}

		const maxTokens = generalConfig.maxTokens || 4096;

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

	private prepareContents(options: IPromptOptions): any[] {
		const contents: any[] = [];

		if (options.messages) {
			for (const msg of options.messages) {
				// Skip system messages (handled separately in Gemini API)
				if (msg.role === 'system') continue;

				// Map OpenAI roles to Gemini roles
				let role: string = msg.role;
				if (role === 'assistant') role = 'model';

				const parts: any[] = [];

				// Handle multi-modal (image input)
				if (msg.files && msg.files.length > 0) {
					parts.push({ text: msg.content });

					for (const file of msg.files) {
						if (file.type?.startsWith('image/')) {
							// Gemini supports inline data or file URIs
							if (file.url.startsWith('data:')) {
								const [mimeType, base64Data] = file.url.split(',');
								const mime = mimeType.split(':')[1].split(';')[0];
								parts.push({
									inlineData: {
										mimeType: mime,
										data: base64Data
									}
								});
							} else {
								parts.push({
									fileData: {
										mimeType: file.type,
										fileUri: file.url
									}
								});
							}
						}
					}
				} else {
					parts.push({ text: msg.content });
				}

				contents.push({ role, parts });
			}
		}

		return contents;
	}

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

		const url = `${this.baseUrl}/${model}:streamGenerateContent?key=${apiKey}&alt=sse`;

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

		// Use native fetch for better control over streaming
		fetch(url, {
			method: 'POST',
			headers: {
				'Content-Type': 'application/json',
			},
			body: JSON.stringify(payload)
		}).then(async (response) => {
			if (!response.ok) {
				const error = await response.text();
				throw new Error(`Google AI API error: ${error}`);
			}

			if (!response.body) {
				throw new Error('No response body');
			}

			const reader = response.body.getReader();
			const decoder = new TextDecoder();
			let buffer = '';

			try {
				while (true) {
					const { done, value } = await reader.read();

					if (done) {
						libx.log.v('GoogleHandler: stream completed');
						break;
					}

					buffer += decoder.decode(value, { stream: true });

					// Process complete lines (ending with \n)
					const lines = buffer.split('\n');
					// Keep the last line in buffer only if it's incomplete (doesn't end with \n)
					buffer = buffer.endsWith('\n') ? '' : (lines.pop() || '');

					for (const line of lines) {
						const trimmed = line.trim();
						if (!trimmed || !trimmed.startsWith('data:')) continue;

						const jsonStr = trimmed.slice(5).trim();
						if (!jsonStr) continue;

						try {
							const data = JSON.parse(jsonStr);
							const candidates = data?.candidates;
							if (candidates && candidates.length > 0) {
								const content = candidates[0]?.content;
								const parts = content?.parts;
								if (parts && parts.length > 0) {
									const text = parts[0]?.text;
									if (text) {
										await this.writeStreamChunk(writer, encoder, text, isPlain);
									}
								}
							}
						} catch (e) {
							// Might be incomplete JSON, will be completed in next chunk
							libx.log.v('GoogleHandler: JSON parse error (might be incomplete)', e);
						}
					}
				}

				// Process any remaining incomplete line in buffer
				if (buffer.trim()) {
					const trimmed = buffer.trim();
					if (trimmed.startsWith('data:')) {
						const jsonStr = trimmed.slice(5).trim();
						if (jsonStr) {
							try {
								const data = JSON.parse(jsonStr);
								const candidates = data?.candidates;
								if (candidates && candidates.length > 0) {
									const content = candidates[0]?.content;
									const parts = content?.parts;
									if (parts && parts.length > 0) {
										const text = parts[0]?.text;
										if (text) {
											await this.writeStreamChunk(writer, encoder, text, isPlain);
										}
									}
								}
							} catch (e) {
								libx.log.w('GoogleHandler: JSON parse remaining buffer error', e);
							}
						}
					}
				}

				// Send completion marker
				if (!isPlain) {
					writer.write(encoder.encode('data: [DONE]\n\n'));
				}
				writer.close();

			} catch (error: any) {
				libx.log.e('GoogleHandler: 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();
			}
		}).catch((error) => {
			libx.log.e('GoogleHandler: fetch 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, model: string, isPlain: boolean): Promise<ReadableStream> {
		const url = `${this.baseUrl}/${model}:generateContent?key=${apiKey}`;

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

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

		const data = await response.json();
		const candidates = data?.candidates;
		let content = '';
		if (candidates && candidates.length > 0) {
			const parts = candidates[0]?.content?.parts;
			if (parts && parts.length > 0) {
				content = parts[0]?.text || '';
			}
		}

		const metadata = {
			model,
			usage: data.usageMetadata,
			finish_reason: candidates?.[0]?.finishReason
		};

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