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

export type Memory = {
	name: string;
	description: string;
	content: string;
	embedding: number[];
}

export type MemoryRef = {
	name: string;
	description: string;
}

export type FactBucket = {
	subject: string;
	facts: string[];
}

export type MemoryNode = {
	name: string;
	missing: boolean;
	links: string[];
	backlinks: string[];
}

function extractLinks(content: string): string[] {
	if(!content) return [];
	const matches = content.matchAll(/\[\[([^\]]+)\]\]/g);
	return [...new Set([...matches].map(m => m[1].trim()))];
}

export function extractMetadata(content: string): {links: string[], backlinks: string[]} {
	const match = content.match(/^---\n([\s\S]*?)\n---/);
	if (!match) return {links: [], backlinks: []};

	const fm = match[1];
	const getList = (key: string): string[] => {
		const m = fm.match(new RegExp(`^${key}:\\s*\\[(.*)\\]$`, 'm'));
		if (!m || !m[1].trim()) return [];
		return m[1].split(',').map(s => s.trim().replace(/^"|"$/g, '')).filter(Boolean);
	};

	return {
		links: getList('links'),
		backlinks: getList('backlinks'),
	};
}

function dedupeFacts(facts: string[]): string[] {
	const seen = new Map<string, string>();
	for (const f of facts) {
		const clean = f.trim();
		if (clean) seen.set(clean.toLowerCase(), clean);
	}
	return [...seen.values()];
}

function cosineDistance(a: number[], b: number[]): number {
	let dot = 0, normA = 0, normB = 0;
	for (let i = 0; i < a.length; i++) {
		dot   += a[i] * b[i];
		normA += a[i] * a[i];
		normB += b[i] * b[i];
	}
	const denom = Math.sqrt(normA) * Math.sqrt(normB);
	return denom === 0 ? 1 : 1 - dot / denom;
}

function getWeekMonday(date: Date = new Date()): string {
	const d = new Date(Date.UTC(date.getFullYear(), date.getMonth(), date.getDate()));
	const day = d.getUTCDay();
	const diff = day === 0 ? -6 : 1 - day;
	d.setUTCDate(d.getUTCDate() + diff);
	return d.toISOString().slice(0, 10);
}

function getWeekSunday(monday: string): string {
	const d = new Date(`${monday}T00:00:00Z`);
	d.setUTCDate(d.getUTCDate() + 6);
	return d.toISOString().slice(0, 10);
}



export class MemoryManager {
	private pendingMemorizations = new Map<string, {
		memories: Memory[] | MemoryCache,
		tempMemoryName: string,
		timestamp: number,
	}>();

	private queues = new Map<string, {
		pending: string[],
		request: {abort?: () => void} | null,
		task: Promise<void>,
	}>();

	tools = {
		read: (memories: Memory[] | MemoryCache): AiTool => ({
			name: 'memory_recall',
			description: 'Read the full content of a memory document',
			args: {
				name: {type: 'string', description: 'Exact memory name', required: true},
			},
			fn: (args: any) => {
				const mems = memories instanceof MemoryCache ? memories.memories : memories;
				const mem = mems.find(m => m.name === args.name);
				if (!mem) return 'Document not found';
				return mem.content;
			},
		}),

		forget: (memories: Memory[] | MemoryCache): AiTool => ({
			name: 'memory_forget',
			description: 'Permanently delete a memory document and clean up all references to it',
			args: {
				name: {type: 'string', description: 'Exact memory name to forget', required: true}
			},
			fn: (args: any) => {
				const result = this.forget(args.name, memories);
				return result ? `Forgotten: ${args.name}` : `Not found: ${args.name}`;
			},
		}),
	};

	constructor(private llm: any) {}

	private async createTempMemory(conversation: string): Promise<Memory> {
		const timestamp = Date.now();
		const content = `---
name: _temp_${timestamp}
description: Temporary memory - processing in background
tags: [_temporary]
links: []
backlinks: []
modified: ${new Date().toISOString()}
---

# Recent Conversation (Processing)

${conversation}`;
		const [e] = await this.llm.embedding(content);
		return {
			name: `_temp_${timestamp}`,
			description: 'Temporary memory - processing in background',
			content,
			embedding: e?.embedding || [],
		};
	}

	forget(name: string, memories: Memory[] | MemoryCache): boolean {
		const mem = memories instanceof MemoryCache ? memories.memories : memories;
		const idx = mem.findIndex(m => m.name === name);
		if (idx === -1) return false;

		for (const node of mem) {
			const {links, backlinks} = extractMetadata(node.content);
			const newBacklinks = backlinks.filter(b => b !== name);
			const newLinks = links.filter(l => l !== name);

			if (newBacklinks.length !== backlinks.length || newLinks.length !== links.length) {
				node.content = this.updateFrontmatter(node.content, {
					links: newLinks,
					backlinks: newBacklinks,
				});
			}
		}

		mem.splice(idx, 1);

		if (memories instanceof MemoryCache) memories.rebuild();
		return true;
	}

	private cosineSearch(query: number[], memories: Memory[], limit: number): MemoryRef[] {
		const scored = memories
			.filter(m => m.embedding?.length)
			.map(m => ({
				ref: {name: m.name, description: m.description},
				distance: cosineDistance(query, m.embedding),
			}))
			.sort((a, b) => a.distance - b.distance)
			.slice(0, limit);
		return scored.map(s => s.ref);
	}

	private listNodes(memories: Memory[]): MemoryRef[] {
		return memories.map(m => ({name: m.name, description: m.description}));
	}

	async recollect(query: string, memories: Memory[] | MemoryCache, limit = 5, graphDepth = 1): Promise<Memory[]> {
		const mem: Memory[] = memories instanceof MemoryCache ? memories.memories : memories;
		if (!mem.length) return [];

		const [e] = await this.llm.embedding(query);
		if (!e) return [];

		let vectorResults: MemoryRef[];
		if (memories instanceof MemoryCache) vectorResults = memories.search(e.embedding, limit);
		else vectorResults = this.cosineSearch(e.embedding, mem, limit);
		const found = new Set<string>(vectorResults.map(r => r.name));

		if (graphDepth > 0) {
			const frontier = [...found];
			for (let depth = 0; depth < graphDepth; depth++) {
				const next: string[] = [];
				for (const name of frontier) {
					const node = mem.find(m => m.name === name);
					if (!node) continue;
					const {links} = extractMetadata(node.content);
					for (const link of links) {
						if (!found.has(link) && mem.find(m => m.name === link)) {
							found.add(link);
							next.push(link);
						}
					}
				}
				frontier.splice(0, frontier.length, ...next);
				if (!frontier.length) break;
			}
		}

		const vectorOrder = vectorResults.map(r => r.name);
		const graphExpansions = [...found].filter(n => !vectorOrder.includes(n));
		const ordered = [...vectorOrder, ...graphExpansions];
		return ordered.map(n => mem.find(m => m.name === n)!).filter(Boolean);
	}

	async memorize(history: LLMMessage[], memories: Memory[] | MemoryCache, options: LLMRequest): Promise<Memory[]> {
		const conversation = history
			.filter(h => h.role === 'user' || h.role === 'assistant')
			.map(h => `[${h.role}]: ${h.content}`).join('\n\n').trim();
		if(!conversation) return [];

		const trackingId = `${Date.now()}_${Math.random()}`;
		const tempMemory = await this.createTempMemory(conversation);
		const mem = memories instanceof MemoryCache ? memories.memories : memories;
		mem.push(tempMemory);
		if (memories instanceof MemoryCache) memories.rebuild();
		this.pendingMemorizations.set(trackingId, {
			memories,
			tempMemoryName: tempMemory.name,
			timestamp: Date.now(),
		});

		try {
			await this._memorizeBackground(conversation, memories, options, tempMemory.name);
			const finalMem = memories instanceof MemoryCache ? memories.memories : memories;
			return finalMem.filter(m => !m.name.startsWith('_temp_'));
		} finally {
			const pending = this.pendingMemorizations.get(trackingId);
			if (pending) {
				const cleanMem = pending.memories instanceof MemoryCache
					? pending.memories.memories
					: pending.memories;
				const idx = cleanMem.findIndex(m => m.name === pending.tempMemoryName);
				if (idx !== -1) cleanMem.splice(idx, 1);
				if (pending.memories instanceof MemoryCache) pending.memories.rebuild();
			}
			this.pendingMemorizations.delete(trackingId);
		}
	}

	private async _memorizeBackground(conversation: string, memories: Memory[] | MemoryCache, options: LLMRequest, tempName: string): Promise<void> {
		const mem = memories instanceof MemoryCache ? memories.memories : memories;
		const monday = getWeekMonday();
		const sunday = getWeekSunday(monday);
		const buckets = await this.factAgent(conversation, mem, options, monday);
		if(!buckets.length) return;
		const jobs = [...buckets].map(({subject, facts}) => {
			let node = mem.find(m => m.name === subject);
			if(!node) {
				node = {name: subject, description: '', content: '', embedding: [],};
				mem.push(node);
			}
			const week = subject.startsWith('Journal/') ? {monday, sunday} : undefined;
			return this.enqueue(node, facts, mem, options, tempName, week);
		});
		await Promise.all(jobs);
	}

	/**
	 * Coalescing queue: if a doc is already compiling, abort the in-flight run, merge its
	 * facts with the new ones and restart. Never blocks a pending update, never drops facts.
	 */
	private enqueue(node: Memory, facts: string[], memories: Memory[] | MemoryCache, options: LLMRequest, tempName: string, week?: {monday: string, sunday: string}): Promise<void> {
		const key = node.name;
		const existing = this.queues.get(key);
		if (existing) {
			existing.pending.push(...facts);
			existing.request?.abort?.();
			return existing.task;
		}

		const entry: {pending: string[], request: {abort?: () => void} | null, task: Promise<void>} = {pending: [...facts], request: null, task: Promise.resolve()};
		this.queues.set(key, entry);
		const m = memories instanceof MemoryCache ? memories.memories : memories;
		entry.task = (async () => {
			while (entry.pending.length) {
				const batch = dedupeFacts(entry.pending.splice(0, entry.pending.length));
				const written = await this.docAgent(node, batch, m, options, tempName, week, entry);
				if (!written) entry.pending.unshift(...batch);
			}
		})().finally(() => {
			this.queues.delete(key);
			if(!this.queues.size && memories instanceof MemoryCache) memories.rebuild();
		});
		return entry.task;
	}

	private buildHeader(node: Memory, week?: {monday: string, sunday: string}, links: string[] = [], backlinks: string[] = []): string {
		const tags = node.name.split('/')[0]?.toLowerCase();
		const lines = [
			'---',
			`name: ${node.name}`,
			`description: ${node.description || ''}`,
			tags ? `tags: [${tags}]` : '',
			links.length ? `links: [${links.map(l => `"${l}"`).join(', ')}]` : 'links: []',
			backlinks.length ? `backlinks: [${backlinks.map(l => `"${l}"`).join(', ')}]` : 'backlinks: []',
			week ? `week: ${week.monday} – ${week.sunday}` : '',
			`modified: ${new Date().toISOString()}`,
			'---',
		].filter(Boolean);
		return lines.join('\n');
	}

	private applyHeader(content: string, header: string): string {
		return `${header}\n\n${this.stripHeader(content)}`;
	}

	private updateFrontmatter(content: string, updates: {links?: string[], backlinks?: string[]}): string {
		const match = content.match(/^---\n([\s\S]*?)\n---\n\n?([\s\S]*)$/);
		if (!match) return content;

		const [, fm, body] = match;
		let newFm = fm;

		if (updates.links !== undefined) {
			const linksList = updates.links.length ? `[${updates.links.map(l => `"${l}"`).join(', ')}]` : '[]';
			newFm = newFm.replace(/^links:.*$/m, `links: ${linksList}`);
		}

		if (updates.backlinks !== undefined) {
			const backlinksList = updates.backlinks.length ? `[${updates.backlinks.map(l => `"${l}"`).join(', ')}]` : '[]';
			newFm = newFm.replace(/^backlinks:.*$/m, `backlinks: ${backlinksList}`);
		}

		newFm = newFm.replace(/^modified:.*$/m, `modified: ${new Date().toISOString()}`);

		return `---\n${newFm}\n---\n\n${body}`;
	}

	private stripHeader(content: string): string {
		return content.replace(/^---[\s\S]*?\n---\n?/, '').trimStart();
	}

	private async docAgent(node: Memory, facts: string[], memories: Memory[], options: LLMRequest, tempName: string, week: {monday: string, sunday: string} | undefined, entry: {request: {abort?: () => void} | null}): Promise<boolean> {
		const {links: oldLinks} = extractMetadata(node.content);
		const currentBody = this.stripHeader(node.content);
		let update;
		try {
			for(let i = 0; i < 3 && !update?.content; i++) {
				const request = this.llm.ask(`New Facts:\n${facts.map(f => `- ${f}`).join('\n')}`, {
					model: options.model,
					temperature: 0.3,
					schema: {
						description: {type: 'string', description: 'One-line description of what this document covers, no formatting or emojis', required: true},
						content: {type: 'string', description: 'Rewritten document in markdown, without the frontmatter block', required: true},
					},
					system: `You are a knowledge base editor. Rewrite the current document below so it incorporates the new facts.

Formatting rules:
- Use Obsidian-style markdown: # headings, **bold** to add emphasis, __italics__ for titles, terms, etc, bullet & numbered lists for grouped 1D data and tables for 2D data
- Link related concepts with [[WikiLink]] notation using full paths like [[People/Sarah]] or [[Projects/Website]]
- Create links for specific entities (person, place, project, program) and abstract concepts (quantum mechanics, entropy) but skip generics (car, red, dog)
- Keep the document concise, factual, and human-readable
- Resolve contradictions: the new facts always win — delete the outdated statement entirely, never keep both
- Later facts in the list override earlier ones
- Do not add frontmatter blocks, filler, preamble, or AI commentary
${week ? '- This is a weekly journal entry.\n' : ''}
All nodes:
${this.listNodes(memories).map(n => n.name).join(', ') || 'none'}

Current document:
\`\`\`markdown
${currentBody}
\`\`\``}
				);
				entry.request = request;
				update = await request;
			}
		} catch (err: any) {
			if (err?.name === 'AbortError') return false;
			throw err;
		} finally {
			entry.request = null;
		}

		if(!update?.content) return false;
		const newLinks = extractLinks(update.content).filter(l => l !== node.name && l !== tempName);
		const newLinkSet = new Set(newLinks);
		const oldLinkSet = new Set(oldLinks);

		for (const added of newLinkSet) {
			if (!oldLinkSet.has(added)) {
				const target = memories.find(m => m.name === added);
				if (target) {
					const {backlinks} = extractMetadata(target.content);
					if (!backlinks.includes(node.name)) {
						target.content = this.updateFrontmatter(target.content, {
							backlinks: [...backlinks, node.name],
						});
					}
				}
			}
		}
		for (const removed of oldLinkSet) {
			if (!newLinkSet.has(removed)) {
				const target = memories.find(m => m.name === removed);
				if (target) {
					const {backlinks} = extractMetadata(target.content);
					target.content = this.updateFrontmatter(target.content, {
						backlinks: backlinks.filter(b => b !== node.name),
					});
				}
			}
		}

		const {backlinks} = extractMetadata(node.content);
		node.description = node.name !== 'Person/User' ? update.description : 'All information about the current user';
		node.content = this.applyHeader(update.content, this.buildHeader(node, week, newLinks, backlinks));
		const [e] = await this.llm.embedding(node.content);
		if(e) node.embedding = e.embedding;
		return true;
	}

	private async factAgent(conversation: string, memories: Memory[], options: LLMRequest, weekKey: string): Promise<FactBucket[]> {
		const buckets = new Map<string, string[]>();
		await this.llm.ask(conversation, {
			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 current facts the USER explicitly stated about themselves, their work, or their projects
- ONLY extract decisions that were MADE during this conversation
- DO NOT extract anything the AI said, its capabilities, or meta-conversation about the AI
- DO NOT extract greetings, pleasantries, or generic exchanges
- DO NOT extract deltas or changes in facts; ONLY the end fact
- If nothing worth remembering was said, do not call any tools

When extracting facts, you MUST also decide the exact destination path:
- Use an existing node name if the facts clearly belong there
- All information primary about the user should go under "People/User"
- When required, create a new path following collection/subject format (e.g., People/Sarah, Projects/Oxide)
- For journal entries, use "Journal"

Available nodes:
- Journal
${this.listNodes(memories).filter(n => !n.name.includes('_temp_') && !n.name.includes('Journal')).map(n => `- ${n.name}: ${n.description}`).join('\n') || 'None yet.'}`,
			tools: [{
				name: 'facts_extract',
				description: 'Submit facts with their destination',
				args: {
					destination: {type: 'string', description: 'Exact existing node name OR new path (e.g. "People/Sarah", "Projects/Oxide")', required: true},
					facts: {type: 'string', description: 'Comma-separated facts', required: true},
				},
				fn: (args: any) => {
					const subject = args.destination.trim().toLowerCase() === 'journal'
						? `Journal/${weekKey}` : args.destination.trim();
					const facts = buckets.get(subject) ?? [];
					facts.push(...dedupeFacts(String(args.facts).split(',')));
					buckets.set(subject, facts);
					return 'Recorded';
				},
			}],
		});
		return buckets.entries().toArray().map(([subject, facts]) => ({subject, facts}));
	}
}