import assert from "node:assert/strict";
import {
	chmodSync,
	copyFileSync,
	mkdirSync,
	mkdtempSync,
	rmSync,
} from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, beforeEach, test } from "vitest";

import {
	clearRtkSessionBaseline,
	ensureRtkSessionBaseline,
	formatRtkSavingsLabel,
	getRtkSessionSavings,
} from "../resources/extensions/shared/rtk-session-stats.ts";
import { createFakeRtk } from "./rtk-test-utils.ts";

// Store original env values for restoration
let originalRtkDisabled: string | undefined;

beforeEach(() => {
	// Save and clear SF_RTK_DISABLED (and SF_RTK_DISABLED) so tests can use fake RTK binaries
	originalRtkDisabled =
		process.env.SF_RTK_DISABLED ?? process.env.SF_RTK_DISABLED;
	delete process.env.SF_RTK_DISABLED;
	delete process.env.SF_RTK_DISABLED;
});

afterEach(() => {
	// Restore original env
	delete process.env.SF_RTK_DISABLED;
	if (originalRtkDisabled !== undefined) {
		process.env.SF_RTK_DISABLED = originalRtkDisabled;
	} else {
		delete process.env.SF_RTK_DISABLED;
	}
});

function summary(
	totalCommands: number,
	totalInput: number,
	totalOutput: number,
	totalSaved: number,
	totalTimeMs = 1000,
) {
	return JSON.stringify({
		summary: {
			total_commands: totalCommands,
			total_input: totalInput,
			total_output: totalOutput,
			total_saved: totalSaved,
			avg_savings_pct: totalInput > 0 ? (totalSaved / totalInput) * 100 : 0,
			total_time_ms: totalTimeMs,
			avg_time_ms: totalCommands > 0 ? totalTimeMs / totalCommands : 0,
		},
	});
}

test("RTK session savings diff from a persisted baseline", () => {
	const basePath = mkdtempSync(join(tmpdir(), "sf-rtk-session-stats-"));
	mkdirSync(join(basePath, ".sf", "runtime"), { recursive: true });

	const first = createFakeRtk({
		"gain --all --format json": { stdout: summary(10, 1000, 600, 400) },
	});
	const second = createFakeRtk({
		"gain --all --format json": { stdout: summary(14, 1600, 900, 700, 1800) },
	});

	const previous = process.env.SF_RTK_PATH;
	try {
		process.env.SF_RTK_PATH = first.path;
		ensureRtkSessionBaseline(basePath, "sess-1");

		process.env.SF_RTK_PATH = second.path;
		const savings = getRtkSessionSavings(basePath, "sess-1");
		assert.ok(savings, "expected RTK savings snapshot");
		assert.equal(savings?.commands, 4);
		assert.equal(savings?.inputTokens, 600);
		assert.equal(savings?.outputTokens, 300);
		assert.equal(savings?.savedTokens, 300);
		assert.equal(Math.round(savings?.savingsPct ?? 0), 50);
	} finally {
		if (previous === undefined) delete process.env.SF_RTK_PATH;
		else process.env.SF_RTK_PATH = previous;
		first.cleanup();
		second.cleanup();
		rmSync(basePath, { recursive: true, force: true });
	}
});

test("RTK session savings baseline resets cleanly when tracking totals go backwards", () => {
	const basePath = mkdtempSync(join(tmpdir(), "sf-rtk-session-reset-"));
	mkdirSync(join(basePath, ".sf", "runtime"), { recursive: true });

	const first = createFakeRtk({
		"gain --all --format json": { stdout: summary(8, 800, 500, 300) },
	});
	const second = createFakeRtk({
		"gain --all --format json": { stdout: summary(1, 100, 80, 20) },
	});

	const previous = process.env.SF_RTK_PATH;
	try {
		process.env.SF_RTK_PATH = first.path;
		ensureRtkSessionBaseline(basePath, "sess-2");

		process.env.SF_RTK_PATH = second.path;
		const savings = getRtkSessionSavings(basePath, "sess-2");
		assert.ok(savings, "expected RTK savings snapshot");
		assert.equal(savings?.commands, 0);
		assert.equal(savings?.savedTokens, 0);
	} finally {
		if (previous === undefined) delete process.env.SF_RTK_PATH;
		else process.env.SF_RTK_PATH = previous;
		first.cleanup();
		second.cleanup();
		rmSync(basePath, { recursive: true, force: true });
	}
});

test("RTK session stats fall back to the managed RTK path when SF_RTK_PATH is unset", () => {
	const basePath = mkdtempSync(join(tmpdir(), "sf-rtk-session-managed-"));
	mkdirSync(join(basePath, ".sf", "runtime"), { recursive: true });

	const fake = createFakeRtk({
		"gain --all --format json": { stdout: summary(6, 900, 500, 400) },
	});
	const managedHome = mkdtempSync(join(tmpdir(), "sf-rtk-home-"));
	const managedDir = join(managedHome, "agent", "bin");
	const managedPath = join(
		managedDir,
		process.platform === "win32" ? "rtk.cmd" : "rtk",
	);
	mkdirSync(managedDir, { recursive: true });
	copyFileSync(fake.path, managedPath);
	if (process.platform !== "win32") {
		chmodSync(managedPath, 0o755);
	}

	const previousHome = process.env.SF_HOME;
	const previousPath = process.env.SF_RTK_PATH;

	try {
		process.env.SF_HOME = managedHome;
		delete process.env.SF_RTK_PATH;

		const env: NodeJS.ProcessEnv = {
			...process.env,
			SF_HOME: managedHome,
		};
		delete env.SF_RTK_PATH;

		const baseline = ensureRtkSessionBaseline(basePath, "sess-managed", env);
		assert.ok(baseline, "expected baseline from managed RTK path");

		const savings = getRtkSessionSavings(basePath, "sess-managed", env);
		assert.ok(savings, "expected savings snapshot from managed RTK path");
		assert.equal(savings?.commands, 0);
	} finally {
		if (previousHome === undefined) delete process.env.SF_HOME;
		else process.env.SF_HOME = previousHome;
		if (previousPath === undefined) delete process.env.SF_RTK_PATH;
		else process.env.SF_RTK_PATH = previousPath;
		fake.cleanup();
		rmSync(managedHome, { recursive: true, force: true });
		rmSync(basePath, { recursive: true, force: true });
	}
});

test("formatRtkSavingsLabel produces a compact footer string", () => {
	assert.equal(
		formatRtkSavingsLabel({
			commands: 5,
			inputTokens: 5949,
			outputTokens: 2905,
			savedTokens: 3044,
			savingsPct: 51.2,
			totalTimeMs: 3200,
			avgTimeMs: 640,
			updatedAt: new Date().toISOString(),
		}),
		"rtk: 3.0k saved (51%)",
	);
	assert.equal(
		formatRtkSavingsLabel({
			commands: 2,
			inputTokens: 0,
			outputTokens: 0,
			savedTokens: 0,
			savingsPct: 0,
			totalTimeMs: 120,
			avgTimeMs: 60,
			updatedAt: new Date().toISOString(),
		}),
		"rtk: active (2 cmds)",
	);
	assert.equal(formatRtkSavingsLabel(null), null);
});

test("clearRtkSessionBaseline removes a stored session entry", () => {
	const basePath = mkdtempSync(join(tmpdir(), "sf-rtk-session-clear-"));
	mkdirSync(join(basePath, ".sf", "runtime"), { recursive: true });
	const fake = createFakeRtk({
		"gain --all --format json": { stdout: summary(3, 300, 200, 100) },
	});
	const previous = process.env.SF_RTK_PATH;

	try {
		process.env.SF_RTK_PATH = fake.path;
		ensureRtkSessionBaseline(basePath, "sess-clear");
		clearRtkSessionBaseline(basePath, "sess-clear");
		const savings = getRtkSessionSavings(basePath, "sess-clear");
		assert.ok(savings, "expected savings snapshot after baseline recreation");
		assert.equal(savings?.commands, 0);
	} finally {
		if (previous === undefined) delete process.env.SF_RTK_PATH;
		else process.env.SF_RTK_PATH = previous;
		fake.cleanup();
		rmSync(basePath, { recursive: true, force: true });
	}
});
