import {OpenAI as openAI} from 'openai';
import {findByProp, objectMap, JSONSanitize, JSONAttemptParse, clean, makeArray} from '@ztimson/utils';
import {AbortablePromise, Ai} from './ai.ts';
import {LLMMessage, LLMRequest} from './llm.ts';
import {LLMProvider} from './provider.ts';
import {TokenPool} from './token-pool.ts';
import {convertSchema} from './tools.ts';

export class OpenAi extends LLMProvider {
	tokenPool!: TokenPool;
	private clients = new Map<string, openAI>();

	constructor(public readonly ai: Ai, public readonly host: string | null, public readonly token: string | string[], public model: string) {
		super();
		const tokens = makeArray(token).filter(Boolean);
		this.tokenPool = new TokenPool(...(tokens.length ? tokens : [host ? 'ignored' : '']));
	}

	private getClient(token: string): openAI {
		let client = this.clients.get(token);
		if(!client) {
			client = new openAI(clean({baseURL: this.host, apiKey: token || undefined}));
			this.clients.set(token, client);
		}
		return client;
	}

	private toWireContent(content: any): any {
		if(!Array.isArray(content)) return content;
		return content.map(c => c.type === 'image'
			? {type: 'image_url', image_url: {url: `data:${c.mime};base64,${c.data}`}}
			: {type: 'text', text: c.text});
	}

	/** Convert standard history -> OpenAI wire format */
	private toWire(history: LLMMessage[], system?: string): any[] {
		const wire: any[] = [];
		if(system) wire.push({role: 'system', content: system});

		for(let i = 0; i < history.length; i++) {
			const h = history[i];

			if(h.role !== 'tool') {
				wire.push({role: h.role, content: this.toWireContent(h.content)});
				continue;
			}

			const calls: any[] = [];
			const results: any[] = [];

			while(i < history.length && history[i].role === 'tool') {
				const tool: any = history[i];

				calls.push({
					id: tool.id,
					type: 'function',
					function: {
						name: tool.name,
						arguments: JSON.stringify(tool.args || {})
					}
				});

				results.push({
					role: 'tool',
					tool_call_id: tool.id,
					content: tool.error || tool.content || ''
				});

				i++;
			}

			wire.push({
				role: 'assistant',
				content: null,
				tool_calls: calls
			});

			wire.push(...results);
			i--;
		}

		return wire;
	}

	ask(message: string, options: LLMRequest = {}): AbortablePromise<string | any> {
		const controller = new AbortController();
		return Object.assign(new Promise<any>(async (res, rej) => {
			if(!options.history) options.history = [];
			const history = options.history;
			if(message) history.push({role: 'user', content: message, timestamp: Date.now()});
			const tools = options.tools || this.ai.options.llm?.tools || [];
			const requestParams: any = {
				model: options.model || this.model,
				stream: !!options.stream,
				max_completion_tokens: options.maxTokens ?? this.ai.options.llm?.maxTokens,
				temperature: options.temperature ?? this.ai.options.llm?.temperature,
				tools: tools.map(t => ({
					type: 'function',
					function: {
						name: t.name,
						description: t.description,
						parameters: {
							type: 'object',
							properties: t.args
								? objectMap(t.args, (key, value) => ({...value, required: undefined}))
								: {},
							required: t.args
								? Object.entries(t.args).filter(t => t[1].required).map(t => t[0])
								: []
						}
					}
				}))
			};

			if(options.schema) {
				const schema = convertSchema(options.schema);
				requestParams.response_format = {
					type: 'json_schema',
					json_schema: {name: 'response', strict: true, schema}
				};
			}
			if(options.stream) requestParams.stream_options = {include_usage: true};

			try {
				let terminal = false;
				let iteration = 0;

				do {
					iteration++;
					requestParams.messages = this.toWire(history.filter(h => h.role !== 'system'), options.system);

					const callStart = Date.now();
					const resp: any = await this.tokenPool.run(token =>
						this.getClient(token).chat.completions.create(requestParams)
					).catch(err => {
						err.message += `\n\nMessages:\n${JSON.stringify(requestParams.messages, null, 2)}`;
						throw err;
					});

					let usage: any;
					let finishReason: string | undefined;
					let msg: any = {content: '', tool_calls: []};
					let streamedChars = 0;

					if(options.stream) {
						let streamCompleted = false;
						try {
							for await (const chunk of resp) {
								if(controller.signal.aborted) break;
								if(chunk.usage) usage = chunk.usage;

								const choice = chunk.choices?.[0];
								if(choice?.finish_reason) finishReason = choice.finish_reason;

								if(choice?.delta?.content) {
									msg.content += choice.delta.content;
									streamedChars += choice.delta.content.length;
									options.stream({text: choice.delta.content});
								}

								if(choice?.delta?.tool_calls) {
									for(const deltaTC of choice.delta.tool_calls) {
										const index = deltaTC.index ?? msg.tool_calls.length;
										let existing = msg.tool_calls.find((tc: any) => tc.index === index);

										if(!existing) {
											existing = {index, id: '', function: {name: '', arguments: ''}};
											msg.tool_calls.push(existing);
										}

										if(deltaTC.id) existing.id = deltaTC.id;
										if(deltaTC.function?.name) existing.function.name = deltaTC.function.name;
										if(deltaTC.function?.arguments) existing.function.arguments += deltaTC.function.arguments;
									}
								}
							}

							streamCompleted = true;
						} catch(err) {
							if(!controller.signal.aborted) throw err;
						}

						if(streamCompleted && !finishReason) finishReason = msg.tool_calls.length ? 'tool_calls' : 'stop';
					} else {
						usage = resp.usage;
						finishReason = resp.choices[0].finish_reason;
						msg = resp.choices[0].message;
					}

					const duration = Date.now() - callStart;
					const tps = usage?.completion_tokens && duration > 0 ? usage.completion_tokens / (duration / 1000) : 0;

					if(finishReason === 'length' && !controller.signal.aborted) {
						if(msg.content?.trim()) history.push({role: 'assistant', content: msg.content.trim(), timestamp: Date.now(), duration, tps});
						throw new Error(`[OpenAI] Response hit token limit before completing`);
					}

					if(!finishReason && !controller.signal.aborted) {
						throw new Error('[OpenAI] Completion ended without a usable response');
					}

					const toolCalls = msg.tool_calls || [];

					if(toolCalls.length && !controller.signal.aborted) {
						if(msg.content?.trim()) history.push({role: 'assistant', content: msg.content.trim(), timestamp: Date.now(), duration, tps});

						const entries = toolCalls.map((tc: any) => {
							const entry: any = {
								role: 'tool',
								id: tc.id,
								name: tc.function.name,
								args: JSONAttemptParse(tc.function.arguments, {}),
								content: undefined,
								timestamp: Date.now()
							};

							history.push(entry);
							return {tc, entry};
						});

						await Promise.all(entries.map(async ({tc, entry}: any) => {
							const tool = tools.find(findByProp('name', tc.function.name));
							if(options.stream) options.stream({tool: tc.function.name});
							if(!tool) return entry.error = 'Tool not found';
							try {
								const toolStream = options.stream && ((chunk: any) => {
									if(chunk.done) return;
									options.stream!(chunk);
								});

								const result = await tool.fn(entry.args, toolStream, this.ai, tc.id);
								entry.content = typeof result === 'object' ? JSONSanitize(result) : result;
							} catch(err: any) {
								entry.error = err?.message || err?.toString() || 'Unknown';
							}
						}));
					} else {
						terminal = true;
						const text = (msg.content || '').trim();
						if(text) history.push({role: 'assistant', content: text, timestamp: Date.now(), duration, tps});
					}
				} while(!terminal && !controller.signal.aborted);

				if(options.stream) options.stream({done: true});
				const turnStart = history.map(h => h.role).lastIndexOf('user');
				const finalContent = history.slice(turnStart + 1).reduce((str, h) => h.role === 'assistant' ? str + (h.content || '') : str, '').trim();
				res(options.schema ? JSONAttemptParse(finalContent, finalContent) : finalContent);
			} catch(err) {
				rej(err);
			}
		}), {abort: () => controller.abort()});
	}
}