// memory.ts
import {LLMRequest, LLMMessage} from './llm.ts';
import {AiTool} from './tools.ts';

/** Background information the AI will be fed as a knowledge document */
export type Memory = {
	/** Memory subject */
	name: string;
	/** Short description of what this document contains - used for RAG retrieval */
	description: string;
	/** Full markdown content of the document */
	content: string;
	/** Embedding vector of the description - used for similarity search */
	embedding: number[];
}

export type MemoryCollection = {
	/** Memory subject */
	name: string;
	/** Short description - required if isNew */
	description?: string;
	/** Extracted facts to merge */
	facts: string[];
}

export class MemoryManager {

	tools = {
		edit: (memory: Memory): AiTool => ({
			name: 'edit_memory',
			description: 'Edit a memory. Omit start/end to append. Pass start only to replace from that line on (Note line 0 = first line of content / line AFTER description). Pass start+end to replace a specific range. start=0 replaces the whole document. Returns updated document',
			args: {
				content: {type: 'string', description: 'New content', required: true},
				start: {type: 'number', description: 'First line to replace (0-indexed, inclusive). Omit to append.'},
				end: {type: 'number', description: 'Last line to replace (0-indexed, inclusive). Omit to replace from start to end of doc.'},
			},
			fn: (args: any) => {
				const lines = memory.content ? memory.content.split('\n') : [];
				const newLines = args.content.split('\n');
				if(args.start === undefined) lines.push(...newLines);
				else if(args.end === undefined) lines.splice(args.start, lines.length - args.start, ...newLines);
				else lines.splice(args.start, args.end - args.start + 1, ...newLines);
				memory.content = lines.join('\n');
				return memory.content;
			}
		}),
		extract: (pools: MemoryCollection[]): AiTool => ({
			name: 'extract_facts',
			description: 'Extract a list of facts to group into a single memory',
			args: {
				name: {type: 'string', description: 'Exact name of an existing memory, or a new name if none fits ([pro]nouns only)', required: true},
				description: {type: 'string', description: 'One sentence description of the memory subject', required: true},
				facts: {type: 'string', description: 'Comma separated list of extracted facts', required: true},
			},
			fn: (args: any) => {
				pools.push({
					name: args.name,
					description: args.description,
					facts: args.facts.split(',').map((f: string) => f.trim()).filter(Boolean),
				});
				return 'Success';
			}}),
		read: (memories: Memory[]): AiTool => ({
			name: 'read_memory',
			description: 'Read entire memory',
			args: {
				name: {type: 'string', description: 'Exact memory name', required: true},
			},
			fn: (args: any) => {
				const mem = memories.find(m => m.name === args.name);
				if(!mem) return 'Document not found';
				return `Name: ${mem.name}\nDescription: ${mem.description}\n\n${mem.content}`;
			}
		}),
	}

	constructor(private llm: any, private model?: string) {}

	/**
	 * Extracts facts from conversation and groups them into individual memories
	 * @param {string} conversation Full conversation formatted as [role]: content
	 * @param {Memory[]} memories The user's memory documents
	 * @param {LLMRequest} options LLM options
	 * @returns {Promise<MemoryCollection[]>} Fact pools grouped by target document
	 */
	private async extract(conversation: string, memories: Memory[], options: LLMRequest): Promise<MemoryCollection[]> {
		const existingDocs = memories.map(m => `Name: ${m.name}\nDescription: ${m.description}`).join('\n\n');
		const pools: MemoryCollection[] = [];
		await this.llm.ask(conversation, {
			model: this.model || options.model,
			temperature: 0.2,
			system: `You are a fact extractor. Analyze this conversation and extract facts worth remembering long term.
Rules:
- ONLY extract facts the USER explicitly stated about themselves or their business
- ONLY extract decisions that were MADE during this conversation
- DO NOT extract anything the AI said, its name, capabilities, or how it introduced itself
- DO NOT extract greetings, pleasantries or generic exchanges
- If nothing worth remembering was said, dont do anything, skip calling tools

For each fact decide whether it belongs in an existing document or needs a new one, then call the \`extract_facts\` tool.

Existing documents:\n${existingDocs || 'None yet.'}`,
			tools: [this.tools.extract(pools)]
		});
		return pools;
	}

	/**
	 * Bot 2 - Editor: merges a pool of facts into a specific document using surgical line-based edits.
	 * Receives full document content and uses read + amend tools to make precise edits.
	 * @param {MemoryCollection} newMem The fact pool to merge
	 * @param {Memory[]} memories The user's memory documents
	 * @param {LLMRequest} options LLM options
	 */
	private async edit(newMem: MemoryCollection, memories: Memory[], options: LLMRequest): Promise<void> {
		const existing = memories.find(m => m.name === newMem.name);
		const mem: Memory = existing || {name: newMem.name, description: newMem.description || '', content: '', embedding: []};
		const isNew = !existing;

		await this.llm.ask(newMem.facts.map(f => `- ${f}`).join('\n'),
			{
				model: this.model || options.model,
				temperature: 0.2,
				system: `You are a document editor. Merge the users list of facts into the following document using the \`edit_memory\` tool; call it as many times as necessary:
\`\`\`
${mem.content}
\`\`\``,
				tools: [this.tools.edit(mem)]
			}
		);

		if(isNew || mem.description !== existing?.description) {
			const e = await this.llm.embedding(mem.description);
			mem.embedding = e?.[0]?.embedding;
		}

		if(isNew) memories.push(mem);
		else {
			const idx = memories.findIndex(m => m.name === newMem.name);
			if(idx >= 0) memories[idx] = mem;
		}
	}

	/**
	 * Find relevant memory documents for a query using description embeddings
	 * @param {string} query The query to search against
	 * @param {Memory[]} memories The user's memory documents
	 * @param {number} limit Max number of results to return
	 * @returns {Promise<Memory[]>} The most relevant memory documents
	 */
	async recollect(query: string, memories: Memory[], limit = 5): Promise<Memory[]> {
		const [e] = await this.llm.embedding(query);
		return memories
			.filter(m => m.embedding?.length)
			.map(m => ({...m, score: this.llm.cosineSimilarity(m.embedding, e.embedding)}))
			.toSorted((a: any, b: any) => b.score - a.score)
			.slice(0, limit);
	}

	/**
	 * Two-stage memory pipeline: classify facts from conversation history then surgically merge them into documents.
	 * Bot 1 (classify) extracts and groups facts cheaply. Bot 2 (edit) runs per-document in parallel with full content access.
	 * @param {LLMMessage[]} history Full conversation history to digest
	 * @param {Memory[]} memories The user's memory documents — mutated in place
	 * @param {LLMRequest} options LLM options
	 */
	async memorize(history: LLMMessage[], memories: Memory[], options: LLMRequest): Promise<void> {
		const conversation = history
			.filter(h => h.role === 'user' || h.role === 'assistant')
			.map(h => `[${h.role}]: ${h.content}`)
			.join('\n\n');
		if(!conversation.trim()) return;
		const pools = await this.extract(conversation, memories, options);
		if(!pools.length) return;
		await Promise.all(pools.map(pool => this.edit(pool, memories, options)));
	}
}