import {Anthropic as anthropic} from '@anthropic-ai/sdk';
import {findByProp, objectMap, JSONSanitize, JSONAttemptParse, 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 Anthropic extends LLMProvider {
	private clients = new Map<string, anthropic>();
	tokenPool!: TokenPool;

	constructor(public readonly ai: Ai, public readonly apiToken: string | string[], public model: string) {
		super();
		this.tokenPool = new TokenPool(...makeArray(apiToken).filter(Boolean));
	}

	private getClient(token: string): anthropic {
		let client = this.clients.get(token);
		if(!client) {
			client = new anthropic({apiKey: token});
			this.clients.set(token, client);
		}
		return client;
	}

	/** Convert standard history -> Anthropic wire format */
	private toWire(history: LLMMessage[]): any[] {
		const wire: any[] = [];
		for(const h of history) {
			if(h.role === 'tool') {
				wire.push(
					{role: 'assistant', content: [{type: 'tool_use', id: h.id, name: h.name, input: h.args}]},
					{role: 'user', content: [{type: 'tool_result', tool_use_id: h.id, is_error: !!h.error, content: h.error || h.content || ''}]}
				);
			} else {
				wire.push({role: h.role, content: h.content});
			}
		}
		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,
				max_tokens: options.maxTokens || this.ai.options.llm?.maxTokens || 4096,
				system: options.system || this.ai.options.llm?.system || '',
				temperature: options.temperature || this.ai.options.llm?.temperature || undefined,
				tools: tools.map(t => ({
					name: t.name,
					description: t.description,
					input_schema: {
						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]) : []
					}
				})),
				stream: !!options.stream,
			};

			if(options.schema) {
				requestParams.output_config = {format: {type: 'json_schema', schema: convertSchema(options.schema)}};
			}

			try {
				let terminal = false;
				do {
					requestParams.messages = this.toWire(history.filter(h => h.role !== 'system'));

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

					let usage: any, content: any[] = [];
					if(options.stream) {
						for await (const chunk of resp) {
							if(controller.signal.aborted) break;
							if(chunk.type === 'content_block_start') {
								if(chunk.content_block.type === 'text') content.push({type: 'text', text: ''});
								else if(chunk.content_block.type === 'tool_use') content.push({type: 'tool_use', id: chunk.content_block.id, name: chunk.content_block.name, input: ''});
							} else if(chunk.type === 'content_block_delta') {
								if(chunk.delta.type === 'text_delta') {
									content.at(-1).text += chunk.delta.text;
									options.stream({text: chunk.delta.text});
								} else if(chunk.delta.type === 'input_json_delta') {
									content.at(-1).input += chunk.delta.partial_json;
								}
							} else if(chunk.type === 'content_block_stop') {
								const last = content.at(-1);
								if(last?.type === 'tool_use') last.input = last.input ? JSONAttemptParse(last.input, {}) : {};
							} else if(chunk.type === 'message_delta') {
								if(chunk.usage) usage = chunk.usage;
							} else if(chunk.type === 'message_stop') {
								break;
							}
						}
					} else {
						usage = resp.usage;
						content = resp.content;
					}
					const duration = Date.now() - callStart;
					const tps = usage?.output_tokens && duration > 0 ? usage.output_tokens / (duration / 1000) : 0;

					const toolCalls = content.filter((c: any) => c.type === 'tool_use');
					if(toolCalls.length && !controller.signal.aborted) {
						const text = content.filter((c: any) => c.type === 'text').map((c: any) => c.text).join('\n\n').trim();
						if(text) history.push({role: 'assistant', content: text, timestamp: Date.now(), duration, tps});

						const entries = toolCalls.map((tc: any) => {
							const entry: any = {role: 'tool', id: tc.id, name: tc.name, args: tc.input, 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.name));
							if(options.stream) options.stream({tool: tc.name});
							if(!tool) { entry.error = 'Tool not found'; return; }
							try {
								const toolStream = options.stream && ((chunk: any) => {
									if(chunk.done) { terminal = true; 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 = content.filter((c: any) => c.type === 'text').map((c: any) => c.text).join('\n\n').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()});
	}
}