/**
 * omp extension: run auto-model-router IN the omp process.
 *
 * The MAIN omp session embeds the router: it binds a free OS-assigned port,
 * publishes it to the shared `$AUTO_MODEL_ROUTER_HOME/embed.port`, and registers the
 * auto-model-router provider. Subagents do NOT bind their own router — they are
 * ephemeral worker processes whose PIDs get recycled, so a per-PID port file
 * is a race. Instead every subagent registers the same shared provider and
 * routes to the main session's single router.
 *
 * The discriminator is `ctx.hasUI`: the main interactive session has a UI,
 * subagents do not.
 *
 * Install by adding this file's absolute path to omp's `extensions:` list:
 *
 *   # ~/.omp/agent/config.yml
 *   extensions:
 *     - /path/to/auto-model-router/omp-extension/router-embed.ts
 *     - /path/to/auto-model-router/omp-extension/router-toast.ts
 */

import { appendFileSync, readFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";

import { ompModelsPath } from "../src/cli/config-cmd.ts";
import { loadConfig } from "../src/config/load.ts";
import { startServer } from "../src/server/http.ts";
import { readRemoteRouter, remoteProviderRegistration } from "./remote-logic.ts";
import { refreshAndRewrite, shouldRefresh } from "../src/cli/refresh.ts";
import type { StartedServer } from "../src/server/http.ts";
import type { RouterConfig } from "../src/config/types.ts";

import type { ExtensionAPI } from "@oh-my-pi/pi-coding-agent";

import { EMBED_DUMMY_API_KEY, EMBED_PROVIDER_ID, buildProviderConfig, deriveAgentdoxScope, deriveWorkspaceOrigin, embedPortPath, modelsYmlPort, probeEmbed, readEmbedPort, resolveEmbedPort, writeEmbedPort } from "./embed-logic.ts";
import { ORIGIN_ENV, SCOPE_ENV } from "../src/context/scope.ts";

// The workspace's scope for the MAIN model. omp builds that handle from
// models.yml before this file loads, so its X-Agentdox-Scope cannot come from
// the provider we register; instead the managed entry names SCOPE_ENV as the
// header's value, omp resolves that from the environment on every request,
// and this module runs inside omp's process — so setting it here reaches
// every turn of this session, main and side roles alike. A workspace that
// derives no scope (no folder name) leaves whatever the shell set.
//
// The workspace's ORIGIN (its git remote, normalised) rides the same way under
// ORIGIN_ENV: the folder name is the scope, the repository is the fingerprint
// a team uses to find the project when folder names collide. Read once — the
// remote does not change during a session — and shared with the providers
// registered below. A workspace outside a repository sets nothing.
const workspaceOrigin = deriveWorkspaceOrigin(process.cwd());
{
	const workspaceScope = deriveAgentdoxScope(process.cwd());
	if (workspaceScope !== "") process.env[SCOPE_ENV] = workspaceScope;
	if (workspaceOrigin !== "") process.env[ORIGIN_ENV] = workspaceOrigin;
}

/** omp's models.yml as text, or "" when it does not exist / cannot be read. */
function readModelsYml(): string {
	try {
		return readFileSync(ompModelsPath(), "utf8");
	} catch {
		return "";
	}
}

/**
 * Appends one line to `$AUTO_MODEL_ROUTER_HOME/embed.log`.
 *
 * A FILE, deliberately: `console.*` from an extension does not reach omp's
 * session log, so the first attempt at this diagnostic left no trace anywhere
 * and the port lifecycle had to be reconstructed from netstat and mtimes. Best
 * effort — a logging failure must never break a session.
 */
function writeEmbedLog(line: string): void {
	try {
		appendFileSync(join(routerHome(), "embed.log"), `${new Date().toISOString()} ${line}\n`, "utf8");
	} catch {
		// Unwritable home: the router still works, we just lose the breadcrumb.
	}
}

/**
 * Stops the router when the PROCESS ends — never when a session does.
 * Idempotent: registered once, however many sessions this process hosts.
 */
let exitHooked = false;
function trackProcessExit(): void {
	if (exitHooked) return;
	exitHooked = true;
	// `exit` cannot await, and does not need to: the OS reclaims the socket.
	// The signal hooks exist so a Ctrl-C releases the port promptly.
	for (const signal of ["SIGINT", "SIGTERM"] as const) {
		process.once(signal, () => {
			void app?.stop().catch(() => {});
		});
	}
}

/**
 * Registers the auto-model-router provider (and its virtual models) into omp's model
 * registry at a specific bound port.
 */
function registerRouterProvider(pi: ExtensionAPI, port: number, cfg: RouterConfig, sessionId: string, subagent: boolean): void {
	// cwd is omp's workspace, which is what the agentdox scope is derived from
	// when none is configured explicitly.
	const providerConfig = buildProviderConfig(port, cfg, process.cwd(), workspaceOrigin);
	const headers: Record<string, string> = {};
	if (providerConfig.harnessId !== undefined && providerConfig.harnessId !== "") {
		headers["X-Omp-Harness"] = providerConfig.harnessId;
	}
	// Per-session scoping: lets the toast surface only this session's decisions
	// even when several omp sessions share one embedded router's ledger.
	if (sessionId !== "") headers["X-Omp-Session"] = sessionId;
	// A session without a UI is a subagent (or a headless run): the router
	// routes its turns under server.subagentProfile.
	if (subagent) headers["X-Omp-Subagent"] = "1";
	// Which agentdox project's shared context this workspace's turns draw on.
	if (providerConfig.agentdoxScope !== undefined && providerConfig.agentdoxScope !== "") {
		headers["X-Agentdox-Scope"] = providerConfig.agentdoxScope;
	}
	// The repository behind that folder, for a front door that keeps a registry.
	if (providerConfig.agentdoxOrigin !== undefined && providerConfig.agentdoxOrigin !== "") {
		headers["X-Agentdox-Origin"] = providerConfig.agentdoxOrigin;
	}
	pi.registerProvider(EMBED_PROVIDER_ID, {
		baseUrl: providerConfig.baseUrl,
		api: "openai-completions",
		apiKey: EMBED_DUMMY_API_KEY,
		...(Object.keys(headers).length > 0 ? { headers } : {}),
		models: providerConfig.models.map((m) => ({
			id: m.id,
			name: m.name,
			api: "openai-completions",
			reasoning: false,
			input: ["text", "image"],
			contextWindow: m.contextWindow,
			maxTokens: m.maxTokens,
			cost: {
				input: m.cost.input,
				output: m.cost.output,
				cacheRead: m.cost.cacheRead,
				cacheWrite: m.cost.cacheWrite,
			},
		})),
	});
}

/**
 * The router bound by THIS PROCESS, and the port it serves.
 *
 * Module scope on purpose: Bun caches the module per process, so when omp loads
 * the extension into a second host (its provider-refresh / reload path does
 * exactly that) these stay visible. A second host then REUSES this router
 * instead of binding another port and orphaning every model handle omp already
 * resolved against the first one.
 */
let app: StartedServer | null = null;
let boundPort: number | null = null;

/** `$AUTO_MODEL_ROUTER_HOME`, tilde-expanded, defaulting to ~/.auto-model-router. */
function routerHome(): string {
	const raw = process.env.AUTO_MODEL_ROUTER_HOME ?? join(homedir(), ".auto-model-router");
	return raw === "~" || raw.startsWith("~/") || raw.startsWith("~\\") ? join(homedir(), raw.slice(1)) : raw;
}

export default function (pi: ExtensionAPI): void {
	pi.setLabel("auto-model-router embed");

	// Shared port file, written only by the main session's router.
	const portFile = embedPortPath(routerHome());
	// Each interactive session binds its OWN router on an ephemeral port, so
	// sessions stay independent: no shared process to contend over, and no
	// session left broken because another one exited. `server.port` from
	// config.yml is deliberately NOT used here — that port belongs to the
	// standalone `serve` daemon, which may legitimately be running alongside.
	// Set AUTO_MODEL_ROUTER_PORT to pin a fixed port on purpose.
	const requestedPort = resolveEmbedPort(process.env.AUTO_MODEL_ROUTER_PORT);
	const cfg = loadConfig({ overrides: { server: { host: "127.0.0.1", port: requestedPort } } });

	pi.on("session_start", async (_event, ctx) => {
		// The omp UI session id tags every request so the toast can scope its
		// notifications to that exact session (see router-toast.ts).
		const sessionId = ctx.sessionManager.getSessionId();

		// Remote mode (`auto-model-router connect`): a router elsewhere is the
		// router. Register it as the provider with its key and bind nothing
		// locally; the other extensions find it through remote.json.
		let remote = readRemoteRouter(routerHome());
		if (remote !== null) {
			// A short-lived key is traded a day ahead of its expiry, and every config
			// re-written, so no session ever starts on a dead key. The remote keeps the
			// old key valid until its own expiry, so this session's main handle (which
			// omp resolved from models.yml before we loaded) is not cut either way.
			if (shouldRefresh(remote)) {
				try {
					const fresh = await refreshAndRewrite({ remote });
					remote = { ...remote, key: fresh.key, refreshToken: fresh.refreshToken, keyExpiresAtMs: fresh.keyExpiresAtMs, refreshExpiresAtMs: fresh.refreshExpiresAtMs };
					writeEmbedLog(`remote credential refreshed; key valid until ${new Date(fresh.keyExpiresAtMs).toISOString()}`);
				} catch (err) {
					writeEmbedLog(`remote credential refresh failed: ${err instanceof Error ? err.message : String(err)}`);
				}
			}
			pi.registerProvider(EMBED_PROVIDER_ID, remoteProviderRegistration(remote, sessionId, !ctx.hasUI, cfg.ledger.fallbackBlend, deriveAgentdoxScope(process.cwd()), workspaceOrigin));
			pi.setLabel(`auto-model-router remote (${remote.url.replace(/^https?:\/\//, "")})`);
			writeEmbedLog(`remote mode url=${remote.url} user=${remote.userId} session=${sessionId}`);
			return;
		}

		// This module is cached per PROCESS, so `app` and `boundPort` are
		// process-global even when omp loads the extension into more than one
		// host. A router already bound in this process is therefore reusable:
		// re-register it for the new session id and return. Never rebind — a
		// second bind would take a different port and orphan every model handle
		// omp already resolved against the first one.
		if (app !== null && boundPort !== null) {
			registerRouterProvider(pi, boundPort, cfg, sessionId, !ctx.hasUI);
			return;
		}

		if (!ctx.hasUI) {
			// Subagents and headless (-p) sessions prefer the main session's
			// shared router: one process, one ledger, one place to inspect.
			// The main writes the port file before spawning subagents.
			const shared = readEmbedPort(portFile);
			if (shared !== null && (await probeEmbed(shared))) {
				registerRouterProvider(pi, shared, cfg, sessionId, !ctx.hasUI);
				return;
			}
			// No live interactive session (headless batch runs, CI, the
			// benchmark harness): fall back to binding a private router so
			// `--model auto-model-router/auto` still resolves. Ephemeral by
			// design — it dies with this process and never writes the shared
			// port file, so it can never hijack another session's subagents.
			const started = startServer(cfg);
			if (started.server.port === undefined) return;
			app = started;
			boundPort = started.server.port;
			registerRouterProvider(pi, boundPort, cfg, sessionId, !ctx.hasUI);
			return;
		}

		// Main interactive session: bind this session's OWN router, then register
		// the provider against the exact bound port. Registration happens only
		// here — never at factory load, where a stale shared port would be
		// captured into omp's model registry and defeat the correct bound URL.

		// A fixed port was asked for (env var) and a live router already answers
		// there: share it rather than failing to bind. Ephemeral ports — the
		// default — never take this path, so sessions stay independent.
		if (requestedPort !== 0 && (await probeEmbed(requestedPort))) {
			writeEmbedPort(portFile, requestedPort);
			registerRouterProvider(pi, requestedPort, cfg, sessionId, !ctx.hasUI);
			pi.setLabel(`auto-model-router embed (shared :${requestedPort})`);
			return;
		}

		// ADOPT THE ADVERTISED PORT. This is the fix for "provider error: Unable
		// to connect" on every real turn while utility calls kept working.
		//
		// omp resolves `modelRoles.default` (auto-model-router/auto) from
		// models.yml during STARTUP — before this extension loads, so before we
		// can bind or register anything. That resolved handle is a SNAPSHOT: a
		// later registerProvider replaces the registry entry but cannot rewrite
		// a handle omp already built. models.yml names the port of the LAST
		// session that wrote it, which after a normal restart is the session the
		// user just closed — a dead socket. Utility calls (title generation,
		// auto-thinking) resolve AFTER our registration and so hit the live port,
		// which is exactly the asymmetry that made this look like a router fault.
		// Measured in the field via embed.log:
		//   embed ready pid=61872 port=54985 models.yml-advertised=50596
		//
		// So bind the port omp already resolved against, whenever nothing holds
		// it. Each session still runs its OWN router — this only chooses which
		// port that router listens on. A live peer holding it means the handle
		// works anyway (that peer serves it), and we fall back to ephemeral.
		const advertised = modelsYmlPort(readModelsYml());
		if (requestedPort === 0 && advertised !== null) cfg.server.port = advertised;

		// Fall back to an ephemeral port when the preferred one is taken, rather
		// than leaving this session with no provider at all.
		let started: StartedServer;
		try {
			started = startServer(cfg);
		} catch {
			cfg.server.port = 0;
			started = startServer(cfg);
		}
		const actualPort = started.server.port;
		if (actualPort === undefined) return;
		app = started;
		boundPort = actualPort;

		// Publish the port; subagents and the toast read it from here.
		writeEmbedPort(portFile, actualPort);

		// models.yml is deliberately NOT written. The port belongs to THIS
		// process and changes every launch, so persisting it into a file omp
		// reads at STARTUP — before this extension loads — makes a dead port
		// authoritative for the next session's `modelRoles.default`, and that
		// handle is a snapshot no later registration can repair. That is the
		// regression this whole class of failure came from. The provider is
		// registered dynamically below instead, which is what worked before the
		// file was ever written.
		//
		// `auto-model-router config --write` still exists for anyone who wants a
		// static block on purpose; the adoption above keeps such a block
		// harmless by binding whatever port it names when that port is free.

		// Register BEFORE any await: everything omp resolves after this point
		// picks up the live URL, so the registration must not sit behind I/O.
		registerRouterProvider(pi, actualPort, cfg, sessionId, !ctx.hasUI);
		pi.setLabel(`auto-model-router embed :${actualPort}`);

		// NO `session_shutdown` teardown. That event is emitted from session
		// DISPOSAL — including omp's provider-refresh / extension-reload path,
		// which runs in a throwaway extension host while the real session keeps
		// going. Because this module is cached per process, such a handler stops
		// the LIVE router: every subsequent turn then fails with Bun's
		// "Unable to connect", while utility calls that resolve after a later
		// rebind still work — the exact asymmetry observed in the field. The
		// router's lifetime is the PROCESS, and the OS reclaims the socket when
		// the process exits.
		trackProcessExit();

		writeEmbedLog(
			`embed ready pid=${process.pid} port=${actualPort}` +
				` models.yml-advertised=${advertised ?? "none"}` +
				` models.yml-untouched` +
				` self-probe=${(await probeEmbed(actualPort)) ? "ok" : "FAILED"}` +
				` session=${sessionId}`,
		);
	});
}