import type { AgentTool } from "@jaeswift/jae-agent-core";
import type { ToolResultMessage } from "@jaeswift/jae-ai";
import { Type } from "@sinclair/typebox";
import { html } from "lit";
import { Brain, BrainCircuit, Trash2 } from "lucide";
import { registerToolRenderer, renderHeader } from "./renderer-registry.js";
import type { ToolRenderer, ToolRenderResult } from "./types.js";

export interface MemoryEntry {
	id: string;
	content: string;
	tags: string[];
	timestamp: string;
}

export interface MemoryStore {
	entries: MemoryEntry[];
}

const DB_NAME = "jae-memory";
const DB_VERSION = 1;
const STORE_NAME = "memories";

let db: IDBDatabase | null = null;

async function openDB(): Promise<IDBDatabase> {
	if (db) return db;
	return new Promise((resolve, reject) => {
		const req = indexedDB.open(DB_NAME, DB_VERSION);
		req.onupgradeneeded = () => {
			req.result.createObjectStore(STORE_NAME, { keyPath: "id" });
		};
		req.onsuccess = () => {
			db = req.result;
			resolve(db);
		};
		req.onerror = () => reject(req.error);
	});
}

export async function memorySave(content: string, tags: string[] = []): Promise<string> {
	const db = await openDB();
	const entry: MemoryEntry = { id: crypto.randomUUID(), content, tags, timestamp: new Date().toISOString() };
	return new Promise((resolve, reject) => {
		const tx = db.transaction(STORE_NAME, "readwrite");
		tx.objectStore(STORE_NAME).put(entry);
		tx.oncomplete = () => resolve(entry.id);
		tx.onerror = () => reject(tx.error);
	});
}

export async function memoryLoad(): Promise<MemoryEntry[]> {
	const db = await openDB();
	return new Promise((resolve, reject) => {
		const tx = db.transaction(STORE_NAME, "readonly");
		const req = tx.objectStore(STORE_NAME).getAll();
		req.onsuccess = () => resolve(req.result || []);
		req.onerror = () => reject(req.error);
	});
}

export async function memorySearch(query: string): Promise<MemoryEntry[]> {
	const all = await memoryLoad();
	const q = query.toLowerCase();
	return all.filter((e) => e.content.toLowerCase().includes(q) || e.tags.some((t) => t.toLowerCase().includes(q)));
}

export async function memoryDelete(id: string): Promise<void> {
	const db = await openDB();
	return new Promise((resolve, reject) => {
		const tx = db.transaction(STORE_NAME, "readwrite");
		tx.objectStore(STORE_NAME).delete(id);
		tx.oncomplete = () => resolve();
		tx.onerror = () => reject(tx.error);
	});
}

// --- Save Memory Tool ---
const saveMemorySchema = Type.Object({
	content: Type.String({ description: "Information to remember" }),
	tags: Type.Optional(Type.Array(Type.String(), { description: "Tags for categorisation" })),
});

export const saveMemoryTool: AgentTool<typeof saveMemorySchema, { id: string; content: string }> = {
	name: "memory_save",
	label: "Save Memory",
	description: "Save a piece of information to long-term memory for future sessions.",
	parameters: saveMemorySchema,
	async execute(toolCallId, params, signal) {
		const id = await memorySave(params.content, params.tags || []);
		return {
			content: [{ type: "text", text: `Memory saved with ID: ${id}` }],
			details: { id, content: params.content },
		};
	},
};

// --- Recall Memory Tool ---
const recallMemorySchema = Type.Object({
	query: Type.String({ description: "Search query to find relevant memories" }),
});

export const recallMemoryTool: AgentTool<typeof recallMemorySchema, { results: MemoryEntry[] }> = {
	name: "memory_recall",
	label: "Recall Memory",
	description: "Search long-term memory for relevant information.",
	parameters: recallMemorySchema,
	async execute(toolCallId, params, signal) {
		const results = await memorySearch(params.query);
		const text =
			results.length === 0
				? `No memories found for: ${params.query}`
				: results.map((r) => `[${r.timestamp.slice(0, 10)}] ${r.content}`).join("\n\n");
		return {
			content: [{ type: "text", text }],
			details: { results },
		};
	},
};

// --- Renderers ---
class SaveMemoryRenderer implements ToolRenderer {
	render(params: any, result: ToolResultMessage<{ id: string; content: string }> | undefined): ToolRenderResult {
		const state = result ? (result.isError ? "error" : "complete") : "inprogress";
		return {
			content: html`
        <div class="flex flex-col gap-2">
          ${renderHeader(state, Brain, `Memory saved`)}
          ${result?.details ? html`<div class="text-xs text-muted-foreground truncate">${result.details.content}</div>` : ""}
        </div>
      `,
			isCustom: false,
		};
	}
}

class RecallMemoryRenderer implements ToolRenderer {
	render(params: any, result: ToolResultMessage<{ results: MemoryEntry[] }> | undefined): ToolRenderResult {
		const state = result ? (result.isError ? "error" : "complete") : "inprogress";
		const results = result?.details?.results || [];
		return {
			content: html`
        <div class="flex flex-col gap-3">
          ${renderHeader(state, BrainCircuit, `Memory recall: ${params?.query || ""}`)}
          ${
					results.length > 0
						? html`
            <div class="flex flex-col gap-1">
              ${results.map(
						(r) => html`
                <div class="text-xs p-2 rounded border border-border">
                  <div class="text-muted-foreground">${r.timestamp.slice(0, 10)}</div>
                  <div>${r.content}</div>
                </div>
              `,
					)}
            </div>
          `
						: ""
				}
        </div>
      `,
			isCustom: false,
		};
	}
}

registerToolRenderer("memory_save", new SaveMemoryRenderer());
registerToolRenderer("memory_recall", new RecallMemoryRenderer());

export function createMemoryTools() {
	return [saveMemoryTool, recallMemoryTool];
}
