import {findByProp, objectMap, JSONSanitize, JSONAttemptParse} from '@ztimson/utils';
import {Ai} from './ai.ts';
import {LLMMessage, LLMRequest} from './llm.ts';
import {AbortablePromise, LLMProvider} from './provider.ts';
import {Ollama as ollama} from 'ollama';

export class Ollama extends LLMProvider {
	client!: ollama;

	constructor(public readonly ai: Ai, public host: string, public model: string) {
		super();
		this.client = new ollama({host});
	}

	private toStandard(history: any[]): LLMMessage[] {
		for(let i = 0; i < history.length; i++) {
			if(history[i].role == 'assistant' && history[i].tool_calls) {
				if(history[i].content) delete history[i].tool_calls;
				else {
					history.splice(i, 1);
					i--;
				}
			} else if(history[i].role == 'tool') {
				const error = history[i].content.startsWith('{"error":');
				history[i] = {role: 'tool', name: history[i].tool_name, args: history[i].args, [error ? 'error' : 'content']: history[i].content, timestamp: history[i].timestamp};
			}
			if(!history[i]?.timestamp) history[i].timestamp = Date.now();
		}
		return history;
	}

	private fromStandard(history: LLMMessage[]): any[] {
		return history.map((h: any) => {
			const {timestamp, ...rest} = h;
			if(h.role != 'tool') return rest;
			return {role: 'tool', tool_name: h.name, content: h.error || h.content}
		});
	}

	ask(message: string, options: LLMRequest = {}): AbortablePromise<LLMMessage[]> {
		const controller = new AbortController();
		const response = new Promise<any>(async (res, rej) => {
			let system = options.system || this.ai.options.system;
			let history = this.fromStandard([...options.history || [], {role: 'user', content: message, timestamp: Date.now()}]);
			if(history[0].roll == 'system') {
				if(!system) system = history.shift();
				else history.shift();
			}
			if(options.compress) history = await this.ai.language.compressHistory(<any>history, options.compress.max, options.compress.min);
			if(options.system) history.unshift({role: 'system', content: system})

			const tools = options.tools || this.ai.options.tools || [];
			const requestParams: any = {
				model: options.model || this.model,
				messages: history,
				stream: !!options.stream,
				signal: controller.signal,
				options: {
					temperature: options.temperature || this.ai.options.temperature || 0.7,
					num_predict: options.max_tokens || this.ai.options.max_tokens || 4096,
				},
				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]) : []
						}
					}
				}))
			}

			let resp: any, isFirstMessage = true;
			do {
				resp = await this.client.chat(requestParams).catch(err => {
					err.message += `\n\nMessages:\n${JSON.stringify(history, null, 2)}`;
					throw err;
				});

				if(options.stream) {
					if(!isFirstMessage) options.stream({text: '\n\n'});
					else isFirstMessage = false;
					resp.message = {role: 'assistant', content: '', tool_calls: []};
					for await (const chunk of resp) {
						if(controller.signal.aborted) break;
						if(chunk.message?.content) {
							resp.message.content += chunk.message.content;
							options.stream({text: chunk.message.content});
						}
						if(chunk.message?.tool_calls) resp.message.tool_calls = chunk.message.tool_calls;
						if(chunk.done) break;
					}
				}

				if(resp.message?.tool_calls?.length && !controller.signal.aborted) {
					history.push(resp.message);
					const results = await Promise.all(resp.message.tool_calls.map(async (toolCall: any) => {
						const tool = tools.find(findByProp('name', toolCall.function.name));
						if(!tool) return {role: 'tool', tool_name: toolCall.function.name, content: '{"error": "Tool not found"}'};
						const args = typeof toolCall.function.arguments === 'string' ? JSONAttemptParse(toolCall.function.arguments, {}) : toolCall.function.arguments;
						try {
							const result = await tool.fn(args, this.ai);
							return {role: 'tool', tool_name: toolCall.function.name, args, content: JSONSanitize(result)};
						} catch (err: any) {
							return {role: 'tool', tool_name: toolCall.function.name, args, content: JSONSanitize({error: err?.message || err?.toString() || 'Unknown'})};
						}
					}));
					history.push(...results);
					requestParams.messages = history;
				}
			} while (!controller.signal.aborted && resp.message?.tool_calls?.length);

			if(options.stream) options.stream({done: true});
			res(this.toStandard([...history, {role: 'assistant', content: resp.message?.content}]));
		});

		return Object.assign(response, {abort: () => controller.abort()});
	}
}