/**
 * agent-swarm.js — Multi-agent group orchestrator with Letta-style group topology.
 *
 * Purpose: manage a named group of PersistentAgents with durable SQLite-backed registry,
 * tag-based routing, and group turn execution. The swarm persists agent registrations
 * across restarts and coordinates via the shared MessageBus.
 *
 * Consumer: SwarmDispatchLayer, createDefaultSwarm factory, and direct multi-agent
 * orchestration flows.
 */

import { mkdirSync } from "node:fs";
import { join } from "node:path";
import { sfRoot } from "../paths.js";
import { getDatabase, openDatabase } from "../sf-db.js";
import { UokCoordinationStore } from "./coordination-store.js";
import { MessageBus } from "./message-bus.js";
import { PersistentAgent } from "./persistent-agent.js";

/**
 * Letta-style group topology manager types.
 *
 * Purpose: enumerate valid routing strategies so callers get a type-safe,
 * discoverable set of options rather than ad-hoc strings.
 *
 * Consumer: AgentSwarm constructor and run() dispatch logic.
 */
export const ManagerType = Object.freeze({
	ROUND_ROBIN: "round_robin",
	SUPERVISOR: "supervisor",
	DYNAMIC: "dynamic",
	SLEEPTIME: "sleeptime",
});

const DEFAULT_MANAGER_TYPE = ManagerType.SUPERVISOR;
const DEFAULT_TERMINATION_TOKEN = "DONE!";
const DEFAULT_MAX_TURNS = 20;

function ensureDbOpen(basePath) {
	const dir = sfRoot(basePath);
	mkdirSync(dir, { recursive: true });
	const dbPath = join(dir, "sf.db");
	if (!getDatabase()) {
		openDatabase(dbPath);
	}
	const db = getDatabase();
	if (!db) throw new Error(`AgentSwarm: failed to open database at ${dbPath}`);
	return db;
}

/**
 * AgentSwarm — Multi-agent group with MessageBus coordination and Letta-style topology.
 *
 * Purpose: manage a named group of PersistentAgents with durable SQLite-backed registry,
 * tag-based routing, and group turn execution. The swarm persists agent registrations
 * across restarts and coordinates via the shared MessageBus.
 *
 * Consumer: SwarmDispatchLayer, createDefaultSwarm factory, and direct multi-agent
 * orchestration flows.
 */
export class AgentSwarm {
	/**
	 * Create or resume a named agent swarm.
	 *
	 * Purpose: initialize swarm identity and reload any previously registered
	 * agents from the durable KV registry so swarm state survives restarts.
	 *
	 * Consumer: orchestration entry points and swarm factory functions.
	 *
	 * @param {string} basePath - project root used to locate `.sf/sf.db`
	 * @param {object} options
	 * @param {string} options.name - stable swarm identity key (required)
	 * @param {string} [options.managerType='supervisor'] - routing topology
	 * @param {string} [options.terminationToken='DONE!'] - token that ends a run loop
	 * @param {number} [options.maxTurns=20] - max turn iterations per run()
	 */
	constructor(basePath, options = {}) {
		const {
			name,
			managerType = DEFAULT_MANAGER_TYPE,
			terminationToken = DEFAULT_TERMINATION_TOKEN,
			maxTurns = DEFAULT_MAX_TURNS,
		} = options;
		if (!name) throw new Error("AgentSwarm: options.name is required");
		if (!Object.values(ManagerType).includes(managerType)) {
			throw new Error(
				`AgentSwarm: invalid managerType "${managerType}". Must be one of: ${Object.values(ManagerType).join(", ")}`,
			);
		}

		this._basePath = basePath;
		this._name = name;
		this._managerType = managerType;
		this._terminationToken = terminationToken;
		this._maxTurns = maxTurns;

		/** @type {Map<string, PersistentAgent>} */
		this._agents = new Map();

		const db = ensureDbOpen(basePath);
		this._store = new UokCoordinationStore(db);
		this._bus = new MessageBus(basePath);

		// Persist initial config
		this._store.set(this._configKey(), {
			name,
			managerType,
			terminationToken,
			maxTurns,
		});

		// Reload any previously registered agents from the durable registry
		const registry = this._store.get(this._registryKey());
		if (registry && Array.isArray(registry.agents)) {
			for (const entry of registry.agents) {
				const agent = new PersistentAgent(basePath, {
					name: entry.name,
					role: entry.role,
					tags: entry.tags ?? [],
				});
				this._agents.set(entry.name, agent);
			}
		}
	}

	// ─── KV Key Helpers ────────────────────────────────────────────────────────

	_registryKey() {
		return `swarm:${this._name}:registry`;
	}

	_configKey() {
		return `swarm:${this._name}:config`;
	}

	_directoryKey() {
		return `swarm:${this._name}:directory`;
	}

	// ─── Agent Registration ────────────────────────────────────────────────────

	/**
	 * Add a PersistentAgent to the swarm and broadcast the updated agent directory.
	 *
	 * Purpose: maintain a live roster of agents, keeping every agent's
	 * agent_directory core block in sync so each agent knows its peers.
	 *
	 * Consumer: swarm setup code and dynamic agent onboarding flows.
	 *
	 * @param {PersistentAgent} agent
	 */
	register(agent) {
		if (!agent || typeof agent.identity !== "object") {
			throw new Error(
				"AgentSwarm.register: argument must be a PersistentAgent instance",
			);
		}
		const { name } = agent.identity;
		this._agents.set(name, agent);

		// Rebuild and persist the directory
		const directory = this._buildDirectory();
		this._store.set(this._directoryKey(), directory);

		// Push updated directory to every registered agent's core block
		const directoryJson = JSON.stringify(directory);
		for (const a of this._agents.values()) {
			a.setCoreBlock("agent_directory", directoryJson);
		}

		this.persist();
	}

	/**
	 * Retrieve a registered agent by name.
	 *
	 * Purpose: give callers direct access to a named swarm member for targeted
	 * operations without exposing the internal agent map.
	 *
	 * Consumer: routing logic, tests, and external coordination flows.
	 *
	 * @param {string} agentName
	 * @returns {PersistentAgent | undefined}
	 */
	get(agentName) {
		return this._agents.get(agentName);
	}

	/**
	 * Return all registered agents.
	 *
	 * Purpose: provide a snapshot of the full swarm roster for iteration or
	 * broadcast operations.
	 *
	 * Consumer: run() loops, broadcast helpers, and topology queries.
	 *
	 * @returns {PersistentAgent[]}
	 */
	getAll() {
		return [...this._agents.values()];
	}

	/**
	 * Return all agents whose identity.role matches the given role.
	 *
	 * Purpose: enable role-based routing so coordinators, workers, scouts,
	 * and reviewers can be selected without hardcoding agent names.
	 *
	 * Consumer: route(), getTopology(), and dynamic dispatch flows.
	 *
	 * @param {string} role
	 * @returns {PersistentAgent[]}
	 */
	getByRole(role) {
		return this.getAll().filter((a) => a.identity.role === role);
	}

	/**
	 * Return agents matching Letta-style tag filters.
	 *
	 * Purpose: support fine-grained tag-based routing (matchAll = AND,
	 * matchSome = OR) so callers can select agents by capability labels.
	 *
	 * Consumer: sendToAgentsMatchingTags and dynamic swarm routing.
	 *
	 * @param {string[]} [matchAll=[]] - agent must have ALL of these tags
	 * @param {string[]} [matchSome=[]] - agent must have AT LEAST ONE of these tags
	 * @returns {PersistentAgent[]}
	 */
	getByTags(matchAll = [], matchSome = []) {
		return this.getAll().filter((a) => {
			const tags = a.identity.tags ?? [];
			const allMatch =
				matchAll.length === 0 || matchAll.every((t) => tags.includes(t));
			const someMatch =
				matchSome.length === 0 || matchSome.some((t) => tags.includes(t));
			return allMatch && someMatch;
		});
	}

	// ─── Messaging ─────────────────────────────────────────────────────────────

	/**
	 * Send a message from one registered agent to another.
	 *
	 * Purpose: route a point-to-point swarm message via the shared MessageBus
	 * so it is durably stored and survives process restarts.
	 *
	 * Consumer: swarm orchestration flows delegating work between agents.
	 *
	 * @param {string} fromName
	 * @param {string} toName
	 * @param {*} body
	 * @param {object} [metadata={}]
	 * @returns {string} message ID
	 */
	send(fromName, toName, body, metadata = {}) {
		return this._bus.send(
			`agent:${fromName}`,
			`agent:${toName}`,
			body,
			metadata,
		);
	}

	/**
	 * Broadcast a message from one agent to all registered agents.
	 *
	 * Purpose: fan-out a notification or task to the full swarm in one call.
	 *
	 * Consumer: coordinator agents distributing context or system-level events.
	 *
	 * @param {string} fromName
	 * @param {*} body
	 * @param {object} [metadata={}]
	 * @returns {string[]} message IDs
	 */
	broadcast(fromName, body, metadata = {}) {
		const recipients = this.getAll()
			.map((a) => a.identity.name)
			.filter((n) => n !== fromName);
		return this._bus.broadcast(
			`agent:${fromName}`,
			recipients.map((n) => `agent:${n}`),
			body,
			metadata,
		);
	}

	/**
	 * Send a message to all agents matching the given tag filters.
	 *
	 * Purpose: target capability-tagged subsets of the swarm without the caller
	 * knowing individual agent names.
	 *
	 * Consumer: dynamic dispatch flows that route by capability rather than identity.
	 *
	 * @param {string} fromName
	 * @param {*} body
	 * @param {string[]} [matchAll=[]]
	 * @param {string[]} [matchSome=[]]
	 * @returns {string[]} message IDs
	 */
	sendToAgentsMatchingTags(fromName, body, matchAll = [], matchSome = []) {
		const targets = this.getByTags(matchAll, matchSome).filter(
			(a) => a.identity.name !== fromName,
		);
		return targets.map((a) => this.send(fromName, a.identity.name, body));
	}

	// ─── Routing ───────────────────────────────────────────────────────────────

	/**
	 * Route a DispatchEnvelope to the appropriate swarm agent by workMode/unitType.
	 *
	 * Purpose: map task routing decisions to the correct role agent so callers
	 * don't need to know the swarm roster or role layout.
	 *
	 * Consumer: SwarmDispatchLayer and autonomous dispatch loops.
	 *
	 * @param {{ unitType?: string, workMode?: string }} envelope
	 * @returns {PersistentAgent | undefined}
	 */
	route(envelope) {
		const workMode =
			envelope.workMode ?? this._deriveWorkMode(envelope.unitType);
		switch (workMode) {
			case "build":
			case "repair":
				return this.getByRole("worker")[0];
			case "research":
				return this.getByRole("scout")[0];
			case "review":
				return this.getByRole("reviewer")[0];
			case "plan":
				return this.getByRole("planner")[0] ?? this.getByRole("coordinator")[0];
			case "verify":
				return this.getByRole("verifier")[0];
			case "document":
				return this.getByRole("scribe")[0];
			case "challenge":
				return this.getByRole("adversary")[0];
			default:
				return this.getByRole("coordinator")[0] ?? this.getAll()[0];
		}
	}

	_deriveWorkMode(unitType) {
		if (!unitType) return undefined;
		const t = unitType.toLowerCase();
		if (t.includes("research") || t.includes("scout")) return "research";
		if (t.includes("review")) return "review";
		if (
			t.includes("audit") ||
			t.includes("validate") ||
			t.includes("gate") ||
			t.includes("uat")
		)
			return "verify";
		if (t.includes("repair") || t.includes("fix")) return "repair";
		if (
			t.includes("build") ||
			t.includes("code") ||
			t.includes("implement") ||
			t.includes("execute")
		)
			return "build";
		if (
			t.includes("plan") ||
			t.includes("slice") ||
			t.includes("milestone") ||
			t.includes("roadmap")
		)
			return "plan";
		if (t.includes("doc") || t.includes("rewrite") || t.includes("scribe"))
			return "document";
		if (t.includes("challenge") || t.includes("adversar")) return "challenge";
		return undefined;
	}

	// ─── Topology ──────────────────────────────────────────────────────────────

	/**
	 * Return a structured view of the swarm by role.
	 *
	 * Purpose: give callers a single snapshot of the swarm topology for
	 * display, routing, and diagnostic purposes without exposing internals.
	 *
	 * Consumer: run() orchestration and swarm health checks.
	 *
	 * @returns {{ coordinator: PersistentAgent[], workers: PersistentAgent[], scouts: PersistentAgent[], reviewers: PersistentAgent[], all: PersistentAgent[] }}
	 */
	getTopology() {
		return {
			coordinator: this.getByRole("coordinator"),
			planners: this.getByRole("planner"),
			workers: this.getByRole("worker"),
			scouts: this.getByRole("scout"),
			reviewers: this.getByRole("reviewer"),
			verifiers: this.getByRole("verifier"),
			scribes: this.getByRole("scribe"),
			adversaries: this.getByRole("adversary"),
			all: this.getAll(),
		};
	}

	// ─── Group Turn Loop ───────────────────────────────────────────────────────

	/**
	 * Execute a group turn loop over the swarm agents.
	 *
	 * Purpose: drive a multi-turn conversation across agents using the configured
	 * topology (round_robin, supervisor, dynamic, sleeptime) and return the full
	 * turn history with termination status.
	 *
	 * When `opts.enableLLM` is true, each agent turn is backed by an actual LLM
	 * call via the agent runner. Otherwise the loop is a pure message router
	 * (useful for tests and lightweight orchestration).
	 *
	 * Consumer: autonomous dispatch, eval harness, and multi-agent task runners.
	 *
	 * @param {string} initialMessage
	 * @param {object} [opts]
	 * @param {string} [opts.from] - agent name to send as
	 * @param {number} [opts.maxTurns] - override swarm default
	 * @param {number} [opts.timeout] - per-turn timeout in ms (default 120000)
	 * @param {boolean} [opts.enableLLM] - run LLM-backed turns (default false)
	 * @returns {Promise<{ turns: Array<{agent: string, message: string}>, terminated: boolean, reason: string }>}
	 */
	async run(initialMessage, opts = {}) {
		const maxTurns = opts.maxTurns ?? this._maxTurns;
		const fromName = opts.from ?? "external";
		const terminationToken = this._terminationToken;
		const enableLLM = opts.enableLLM ?? false;

		const turns = [];
		let terminated = false;
		let reason = "max_turns_reached";

		const agents = this.getAll();
		if (agents.length === 0) {
			return { turns, terminated: false, reason: "no_agents" };
		}

		/** @type {((agent: import("./persistent-agent.js").PersistentAgent, opts?: object) => Promise<{response?: string}>) | null} */
		let runAgentTurn = null;
		if (enableLLM) {
			const runner = await import("./agent-runner.js");
			runAgentTurn = runner.runAgentTurn;
		}

		switch (this._managerType) {
			case ManagerType.ROUND_ROBIN: {
				let message = initialMessage;
				for (let i = 0; i < maxTurns; i++) {
					const agent = agents[i % agents.length];
					const agentName = agent.identity.name;
					this.send(fromName, agentName, message);

					let reply;
					if (runAgentTurn) {
						const result = await runAgentTurn(agent, {
							timeoutMs: opts.timeout,
						});
						reply = result.response ?? message;
					} else {
						const received = agent.receive(false);
						const last = received[received.length - 1];
						reply = _bodyToString(last?.body, message);
					}
					turns.push({ agent: agentName, message: reply });
					if (reply.includes(terminationToken)) {
						terminated = true;
						reason = "termination_token";
						break;
					}
					message = reply;
				}
				break;
			}

			case ManagerType.SUPERVISOR: {
				const { coordinator } = this.getTopology();
				const supervisor = coordinator[0] ?? agents[0];
				const supervisorName = supervisor.identity.name;
				this.send(fromName, supervisorName, initialMessage);

				let message = initialMessage;
				for (let i = 0; i < maxTurns; i++) {
					let reply;
					if (runAgentTurn) {
						const result = await runAgentTurn(supervisor, {
							timeoutMs: opts.timeout,
						});
						reply = result.response ?? message;
					} else {
						const received = supervisor.receive(true);
						const last = received[received.length - 1];
						reply = _bodyToString(last?.body, message);
					}
					turns.push({ agent: supervisorName, message: reply });
					if (reply.includes(terminationToken)) {
						terminated = true;
						reason = "termination_token";
						break;
					}

					// Supervisor routes to a worker (cycle through workers)
					const workers = this.getByRole("worker");
					if (workers.length > 0) {
						const worker = workers[i % workers.length];
						const workerName = worker.identity.name;
						this.send(supervisorName, workerName, reply);

						let workerReply;
						if (runAgentTurn) {
							const result = await runAgentTurn(worker, {
								timeoutMs: opts.timeout,
							});
							workerReply = result.response ?? reply;
						} else {
							const workerMessages = worker.receive(true);
							const workerLast = workerMessages[workerMessages.length - 1];
							workerReply = _bodyToString(workerLast?.body, reply);
						}
						turns.push({ agent: workerName, message: workerReply });
						if (workerReply.includes(terminationToken)) {
							terminated = true;
							reason = "termination_token";
							break;
						}
						// Route worker result back to supervisor for next iteration
						this.send(workerName, supervisorName, workerReply);
						message = workerReply;
					} else {
						message = reply;
					}
				}
				break;
			}

			case ManagerType.DYNAMIC: {
				// Each turn: give every agent a chance to reply; first with token ends loop
				let message = initialMessage;
				for (let i = 0; i < maxTurns; i++) {
					let tokenFound = false;
					for (const agent of agents) {
						const agentName = agent.identity.name;
						this.send(fromName, agentName, message);

						let reply;
						if (runAgentTurn) {
							const result = await runAgentTurn(agent, {
								timeoutMs: opts.timeout,
							});
							reply = result.response ?? message;
						} else {
							const received = agent.receive(false);
							const last = received[received.length - 1];
							reply = _bodyToString(last?.body, message);
						}
						turns.push({ agent: agentName, message: reply });
						if (reply.includes(terminationToken)) {
							terminated = true;
							reason = "termination_token";
							tokenFound = true;
							break;
						}
						message = reply;
					}
					if (tokenFound) break;
				}
				break;
			}

			case ManagerType.SLEEPTIME: {
				// One conversation agent + non-blocking memory consolidation enqueued to DB.
				// The scheduler drains sleeptime_consolidation_queue on each poll tick.
				const conversationAgents = agents.filter(
					(a) => a.identity.role !== "coordinator",
				);
				const memoryAgent = this.getByRole("coordinator")[0];
				const convoAgent = conversationAgents[0] ?? agents[0];
				const convoName = convoAgent.identity.name;

				let message = initialMessage;
				this.send(fromName, convoName, message);

				for (let i = 0; i < maxTurns; i++) {
					let reply;
					if (runAgentTurn) {
						const result = await runAgentTurn(convoAgent, {
							timeoutMs: opts.timeout,
						});
						reply = result.response ?? message;
					} else {
						const received = convoAgent.receive(true);
						const last = received[received.length - 1];
						reply = _bodyToString(last?.body, message);
					}
					turns.push({ agent: convoName, message: reply });
					if (reply.includes(terminationToken)) {
						terminated = true;
						reason = "termination_token";
						break;
					}
					message = reply;
					// Enqueue consolidation to DB — non-blocking; scheduler drains it.
					if (memoryAgent) {
						const db = getDatabase();
						if (db) {
							const jobId = `slp-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
							db.prepare(
								`INSERT INTO sleeptime_consolidation_queue
							 (id, conversation_agent, memory_agent, content, status, created_at)
							 VALUES (:id, :convo, :mem, :content, 'pending', :ts)`,
							).run({
								":id": jobId,
								":convo": convoName,
								":mem": memoryAgent.identity.name,
								":content": `consolidate: ${reply}`,
								":ts": new Date().toISOString(),
							});
							turns.push({
								agent: memoryAgent.identity.name,
								message: `[memory:queued] job ${jobId}`,
							});
						}
					}
				}
				break;
			}
		}

		return { turns, terminated, reason };
	}

	// ─── Persistence ───────────────────────────────────────────────────────────

	/**
	 * Save the current agent registry snapshot to the KV store.
	 *
	 * Purpose: make swarm membership durable so agents can be rehydrated on
	 * restart without re-registering them programmatically.
	 *
	 * Consumer: register() and any code that mutates swarm membership.
	 */
	persist() {
		const agents = this.getAll().map((a) => ({
			name: a.identity.name,
			role: a.identity.role,
			tags: a.identity.tags ?? [],
		}));
		this._store.set(this._registryKey(), {
			name: this._name,
			managerType: this._managerType,
			agents,
			config: {
				terminationToken: this._terminationToken,
				maxTurns: this._maxTurns,
			},
		});
	}

	/**
	 * Rehydrate an AgentSwarm from a previously persisted registry.
	 *
	 * Purpose: restore swarm state across process restarts without requiring
	 * callers to re-register agents explicitly.
	 *
	 * Consumer: orchestration bootstrap and restart recovery paths.
	 *
	 * @param {string} basePath
	 * @param {string} name
	 * @returns {AgentSwarm}
	 */
	static load(basePath, name) {
		const db = ensureDbOpen(basePath);
		const store = new UokCoordinationStore(db);
		const registryKey = `swarm:${name}:registry`;
		const saved = store.get(registryKey);

		if (!saved) {
			return new AgentSwarm(basePath, { name });
		}

		const swarm = new AgentSwarm(basePath, {
			name: saved.name ?? name,
			managerType: saved.managerType ?? DEFAULT_MANAGER_TYPE,
			terminationToken:
				saved.config?.terminationToken ?? DEFAULT_TERMINATION_TOKEN,
			maxTurns: saved.config?.maxTurns ?? DEFAULT_MAX_TURNS,
		});

		// Agents are reloaded from the registry inside the constructor.
		// Re-call register() on each to ensure agent_directory blocks are current.
		for (const agent of swarm.getAll()) {
			swarm.register(agent);
		}

		return swarm;
	}

	// ─── Internal Helpers ──────────────────────────────────────────────────────

	_buildDirectory() {
		return this.getAll().map((a) => ({
			name: a.identity.name,
			role: a.identity.role,
			tags: a.identity.tags ?? [],
			agentId: a.identity.agentId,
		}));
	}
}

function _bodyToString(body, fallback) {
	if (body == null) return fallback;
	return typeof body === "string" ? body : JSON.stringify(body);
}
