import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { $flag, $which, logger } from "@gajae-code/utils";
import { TOML } from "bun";
import { spawnOwnedProcess } from "../runtime/process-lifecycle";
import { isProjectControlledPath } from "./path-trust";

/**
 * lspmux integration for LSP server multiplexing.
 *
 * When lspmux is available and running, this module wraps supported LSP server
 * commands to use lspmux client mode, enabling server instance sharing across
 * multiple editor windows.
 *
 * Integration is transparent: if lspmux is unavailable, falls back to direct spawning.
 */

// =============================================================================
// Types
// =============================================================================

interface LspmuxConfig {
	instance_timeout?: number;
	gc_interval?: number;
	listen?: [string, number] | string;
	connect?: [string, number] | string;
	log_filters?: string;
	pass_environment?: string[];
}

interface LspmuxState {
	available: boolean;
	running: boolean;
	binaryPath: string | null;
	config: LspmuxConfig | null;
}

// =============================================================================
// Constants
// =============================================================================

/**
 * Servers that benefit from lspmux multiplexing.
 *
 * lspmux can multiplex any LSP server, but it's most beneficial for servers
 * with high startup cost or significant memory usage.
 */
const DEFAULT_SUPPORTED_SERVERS = new Set([
	"rust-analyzer",
	// Other servers can be added after testing with lspmux
]);

/** Timeout for liveness check (ms) */
const LIVENESS_TIMEOUT_MS = 1000;

/** Cache duration for lspmux state (5 minutes) */
const STATE_CACHE_TTL_MS = 5 * 60 * 1000;

// =============================================================================
// Config Path
// =============================================================================

/**
 * Get the lspmux config path based on platform.
 * Matches Rust's `dirs::config_dir()` behavior.
 */
function getConfigPath(): string {
	const home = os.homedir();
	switch (os.platform()) {
		case "win32":
			return path.join(Bun.env.APPDATA ?? path.join(home, "AppData", "Roaming"), "lspmux", "config.toml");
		case "darwin":
			return path.join(home, "Library", "Application Support", "lspmux", "config.toml");
		default:
			return path.join(Bun.env.XDG_CONFIG_HOME ?? path.join(home, ".config"), "lspmux", "config.toml");
	}
}

// =============================================================================
// State Management
// =============================================================================

let cachedState: LspmuxState | null = null;
let cacheTimestamp = 0;
let cacheCwd: string | null = null;

/**
 * Parse lspmux config.toml file.
 */
async function parseConfig(): Promise<LspmuxConfig | null> {
	try {
		const file = Bun.file(getConfigPath());
		if (!(await file.exists())) {
			return null;
		}
		return TOML.parse(await file.text()) as LspmuxConfig;
	} catch {
		return null;
	}
}

/**
 * Check if lspmux server is running via `lspmux status`.
 */
async function checkServerRunning(binaryPath: string): Promise<boolean> {
	try {
		const owner = spawnOwnedProcess([binaryPath, "status"], {
			stdin: "ignore",
			name: "lspmux:status",
		});
		const proc = owner.child;
		drainStream(proc.stdout);
		drainStream(proc.stderr);

		let timer: NodeJS.Timeout | undefined;
		const timeout = new Promise<null>(resolve => {
			timer = setTimeout(() => resolve(null), LIVENESS_TIMEOUT_MS);
		});
		const exited = await Promise.race([proc.exited, timeout]);
		if (timer) clearTimeout(timer);

		if (exited === null) {
			await owner.dispose();
			await owner.awaitExit({ timeoutMs: 1_000 });
			return false;
		}

		return exited === 0;
	} catch {
		return false;
	}
}

function drainStream(stream: ReadableStream<Uint8Array> | null | undefined): void {
	if (!stream) return;
	void (async () => {
		try {
			const reader = stream.getReader();
			try {
				while (!(await reader.read()).done) {
					// drain only
				}
			} finally {
				reader.releaseLock();
			}
		} catch {
			// Process stream closed or already consumed.
		}
	})();
}

function resolveTrustedLspmuxBinary(cwd: string): string | null {
	const discoveredPath = $which("lspmux");
	if (!discoveredPath) return null;
	if (isProjectControlledPath(discoveredPath, cwd)) return null;

	try {
		return fs.realpathSync(discoveredPath);
	} catch {
		return null;
	}
}

/**
 * Detect lspmux availability and state.
 * Results are cached for STATE_CACHE_TTL_MS.
 *
 * Set GJC_DISABLE_LSPMUX=1 or PI_DISABLE_LSPMUX=1 to disable.
 */
export async function detectLspmux(cwd = process.cwd()): Promise<LspmuxState> {
	const now = Date.now();
	if ($flag("GJC_DISABLE_LSPMUX") || $flag("PI_DISABLE_LSPMUX")) {
		cachedState = { available: false, running: false, binaryPath: null, config: null };
		cacheTimestamp = now;
		cacheCwd = cwd;
		return cachedState;
	}

	if (cachedState && cacheCwd === cwd && now - cacheTimestamp < STATE_CACHE_TTL_MS) {
		return cachedState;
	}

	const binaryPath = resolveTrustedLspmuxBinary(cwd);
	if (!binaryPath) {
		cachedState = { available: false, running: false, binaryPath: null, config: null };
		cacheTimestamp = now;
		cacheCwd = cwd;
		return cachedState;
	}

	const [config, running] = await Promise.all([parseConfig(), checkServerRunning(binaryPath)]);

	cachedState = { available: true, running, binaryPath, config };
	cacheTimestamp = now;
	cacheCwd = cwd;

	if (running) {
		logger.debug("lspmux detected and running", { binaryPath });
	}

	return cachedState;
}

/** Reset cached lspmux detection state for focused tests. */
export function resetLspmuxStateForTesting(): void {
	cachedState = null;
	cacheTimestamp = 0;
	cacheCwd = null;
}

// =============================================================================
// Command Wrapping
// =============================================================================

/**
 * Check if a server command is supported by lspmux.
 */
export function isLspmuxSupported(command: string): boolean {
	// Extract base command name (handle full paths)
	const baseName = command.split("/").pop() ?? command;
	return DEFAULT_SUPPORTED_SERVERS.has(baseName);
}

export interface LspmuxWrappedCommand {
	command: string;
	args: string[];
	env?: Record<string, string>;
}

/**
 * Wrap a server command to use lspmux client mode.
 *
 * @param originalCommand - The original LSP server command (e.g., "rust-analyzer")
 * @param originalArgs - Original command arguments
 * @param state - lspmux state from detectLspmux()
 * @returns Wrapped command, args, and env vars; or original if lspmux unavailable
 */
export function wrapWithLspmux(
	originalCommand: string,
	originalArgs: string[] | undefined,
	state: LspmuxState,
): LspmuxWrappedCommand {
	if (!state.available || !state.running || !state.binaryPath) {
		return { command: originalCommand, args: originalArgs ?? [] };
	}

	if (!isLspmuxSupported(originalCommand)) {
		return { command: originalCommand, args: originalArgs ?? [] };
	}

	const baseName = originalCommand.split("/").pop() ?? originalCommand;
	const isDefaultRustAnalyzer = baseName === "rust-analyzer" && originalCommand === "rust-analyzer";
	const hasArgs = originalArgs && originalArgs.length > 0;

	// rust-analyzer from $PATH with no args - lspmux's default, simplest case
	if (isDefaultRustAnalyzer && !hasArgs) {
		return { command: state.binaryPath, args: [] };
	}

	// Use explicit `client` subcommand with LSPMUX_SERVER env var
	// Use `--` to separate lspmux options from server args
	const args = hasArgs ? ["client", "--", ...originalArgs] : ["client"];
	return {
		command: state.binaryPath,
		args,
		env: { LSPMUX_SERVER: originalCommand },
	};
}

/**
 * Get lspmux-wrapped command if available, otherwise return original.
 * This is the main entry point for config.ts integration.
 *
 * @param command - Original LSP server command
 * @param args - Original command arguments
 * @returns Command and args to use (possibly wrapped with lspmux)
 */
export async function getLspmuxCommand(
	command: string,
	args?: string[],
	cwd = process.cwd(),
): Promise<LspmuxWrappedCommand> {
	const state = await detectLspmux(cwd);
	return wrapWithLspmux(command, args, state);
}