import { describe, it, expect, beforeEach, afterEach } from "vitest";
import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import Database from "better-sqlite3";
import { IpfsService } from "../ipfs.js";
import { RepoManager } from "../repo-manager.js";
import type { Config } from "../config.js";
import { readCarWithRoot } from "@atproto/repo";
import { generateMstProof, verifyMstProof, extractAllRecordPaths, extractAllCids } from "./mst-proof.js";

function testConfig(dataDir: string): Config {
	return {
		DID: "did:plc:test123",
		HANDLE: "test.example.com",
		PDS_HOSTNAME: "test.example.com",
		AUTH_TOKEN: "test-auth-token",
		SIGNING_KEY:
			"0000000000000000000000000000000000000000000000000000000000000001",
		SIGNING_KEY_PUBLIC:
			"zQ3shP2mWsZYWgvZM9GJ3EvMfRXQJwuTh6BdXLvJB9gFhT3Lr",
		JWT_SECRET: "test-jwt-secret",
		PASSWORD_HASH: "$2a$10$test",
		DATA_DIR: dataDir,
		PORT: 3000,
		IPFS_ENABLED: true,
		IPFS_NETWORKING: false,
		REPLICATE_DIDS: [],
		FIREHOSE_URL: "wss://localhost/xrpc/com.atproto.sync.subscribeRepos",
		FIREHOSE_ENABLED: false,
		RATE_LIMIT_ENABLED: false,
		RATE_LIMIT_READ_PER_MIN: 300,
		RATE_LIMIT_SYNC_PER_MIN: 30,
		RATE_LIMIT_SESSION_PER_MIN: 10,
		RATE_LIMIT_WRITE_PER_MIN: 200,
		RATE_LIMIT_CHALLENGE_PER_MIN: 20,
		RATE_LIMIT_MAX_CONNECTIONS: 100,
		RATE_LIMIT_FIREHOSE_PER_IP: 3,
		OAUTH_ENABLED: false, PUBLIC_URL: "http://localhost:3000",
	};
}

describe("MST Path Proof", () => {
	let tmpDir: string;
	let db: InstanceType<typeof Database>;
	let ipfsService: IpfsService;
	let repoManager: RepoManager;

	beforeEach(async () => {
		tmpDir = mkdtempSync(join(tmpdir(), "mst-proof-test-"));
		const config = testConfig(tmpDir);

		db = new Database(join(tmpDir, "test.db"));
		ipfsService = new IpfsService({
			db,
			networking: false,
		});
		await ipfsService.start();

		repoManager = new RepoManager(db, config);
		repoManager.init(undefined, ipfsService, ipfsService);
	});

	afterEach(async () => {
		if (ipfsService.isRunning()) {
			await ipfsService.stop();
		}
		db.close();
		rmSync(tmpDir, { recursive: true, force: true });
	});

	/**
	 * Helper: create records, export CAR, store blocks in IPFS, return root CID.
	 */
	async function getRepoRootCid(): Promise<string> {
		const carBytes = await repoManager.getRepoCar();
		const { root, blocks } = await readCarWithRoot(carBytes);
		await ipfsService.putBlocks(blocks);
		return root.toString();
	}

	// ============================================
	// Existence proofs
	// ============================================

	it("generates and verifies an existence proof for a single record", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Hello, world!",
			createdAt: "2025-01-01T00:00:00.000Z",
		});

		const rootCid = await getRepoRootCid();

		// Get the record's rkey
		const records = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 10,
		});
		const rkey = records.records[0]!.uri.split("/").pop()!;
		const recordPath = `app.bsky.feed.post/${rkey}`;

		// Generate proof
		const proof = await generateMstProof(ipfsService, rootCid, recordPath);

		expect(proof.found).toBe(true);
		expect(proof.recordCid).not.toBeNull();
		expect(proof.commitBlock.cid).toBe(rootCid);
		expect(proof.nodes.length).toBeGreaterThan(0);

		// Verify proof
		const verification = await verifyMstProof(proof, rootCid, recordPath);

		expect(verification.valid).toBe(true);
		expect(verification.found).toBe(true);
		expect(verification.recordCid).toBe(proof.recordCid);
		expect(verification.error).toBeUndefined();
	});

	it("generates and verifies proof with multiple records", async () => {
		// Create several records to build a deeper MST
		for (let i = 0; i < 10; i++) {
			await repoManager.createRecord("app.bsky.feed.post", undefined, {
				$type: "app.bsky.feed.post",
				text: `Post number ${i}`,
				createdAt: new Date().toISOString(),
			});
		}

		const rootCid = await getRepoRootCid();

		// Get all records and verify proofs for each
		const records = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 100,
		});

		for (const record of records.records) {
			const rkey = record.uri.split("/").pop()!;
			const recordPath = `app.bsky.feed.post/${rkey}`;

			const proof = await generateMstProof(
				ipfsService,
				rootCid,
				recordPath,
			);

			expect(proof.found).toBe(true);
			expect(proof.recordCid).not.toBeNull();

			const verification = await verifyMstProof(
				proof,
				rootCid,
				recordPath,
			);
			expect(verification.valid).toBe(true);
			expect(verification.found).toBe(true);
		}
	});

	it("generates and verifies proof for records in different collections", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "A post",
			createdAt: new Date().toISOString(),
		});

		await repoManager.putRecord("app.bsky.actor.profile", "self", {
			$type: "app.bsky.actor.profile",
			displayName: "Test User",
		});

		const rootCid = await getRepoRootCid();

		// Verify post
		const posts = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 10,
		});
		const postRkey = posts.records[0]!.uri.split("/").pop()!;
		const postPath = `app.bsky.feed.post/${postRkey}`;

		const postProof = await generateMstProof(
			ipfsService,
			rootCid,
			postPath,
		);
		expect(postProof.found).toBe(true);
		const postVerification = await verifyMstProof(
			postProof,
			rootCid,
			postPath,
		);
		expect(postVerification.valid).toBe(true);
		expect(postVerification.found).toBe(true);

		// Verify profile
		const profilePath = "app.bsky.actor.profile/self";
		const profileProof = await generateMstProof(
			ipfsService,
			rootCid,
			profilePath,
		);
		expect(profileProof.found).toBe(true);
		const profileVerification = await verifyMstProof(
			profileProof,
			rootCid,
			profilePath,
		);
		expect(profileVerification.valid).toBe(true);
		expect(profileVerification.found).toBe(true);
	});

	// ============================================
	// Non-existence proofs
	// ============================================

	it("generates and verifies a non-existence proof", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Only post",
			createdAt: new Date().toISOString(),
		});

		const rootCid = await getRepoRootCid();

		// Use a path that definitely does not exist
		const nonExistentPath = "app.bsky.feed.post/nonexistent-rkey-12345";

		const proof = await generateMstProof(
			ipfsService,
			rootCid,
			nonExistentPath,
		);

		expect(proof.found).toBe(false);
		expect(proof.recordCid).toBeNull();

		const verification = await verifyMstProof(
			proof,
			rootCid,
			nonExistentPath,
		);

		expect(verification.valid).toBe(true);
		expect(verification.found).toBe(false);
		expect(verification.recordCid).toBeNull();
	});

	it("non-existence proof for nonexistent collection", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Post in a real collection",
			createdAt: new Date().toISOString(),
		});

		const rootCid = await getRepoRootCid();

		const nonExistentPath = "com.example.nonexistent/abc123";

		const proof = await generateMstProof(
			ipfsService,
			rootCid,
			nonExistentPath,
		);

		expect(proof.found).toBe(false);
		expect(proof.recordCid).toBeNull();

		const verification = await verifyMstProof(
			proof,
			rootCid,
			nonExistentPath,
		);
		expect(verification.valid).toBe(true);
		expect(verification.found).toBe(false);
	});

	// ============================================
	// Proof compactness
	// ============================================

	it("proof is compact: fewer blocks than total MST", async () => {
		// Create enough records to build a multi-level MST
		for (let i = 0; i < 20; i++) {
			await repoManager.createRecord("app.bsky.feed.post", undefined, {
				$type: "app.bsky.feed.post",
				text: `Post ${i} for compactness test`,
				createdAt: new Date().toISOString(),
			});
		}

		const rootCid = await getRepoRootCid();

		const records = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 100,
		});
		const rkey = records.records[0]!.uri.split("/").pop()!;
		const recordPath = `app.bsky.feed.post/${rkey}`;

		const proof = await generateMstProof(
			ipfsService,
			rootCid,
			recordPath,
		);

		// The proof should have:
		// - 1 commit block
		// - N MST node blocks (path from root to leaf)
		// For a tree with 20+ entries, this should be fewer blocks than the total tree
		const totalProofBlocks = 1 + proof.nodes.length; // commit + MST nodes
		// 20 records => several MST nodes in total; proof should be a subset
		expect(totalProofBlocks).toBeLessThanOrEqual(10);
		expect(proof.found).toBe(true);
	});

	// ============================================
	// Verification failure cases
	// ============================================

	it("verification fails with wrong commit CID", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Wrong CID test",
			createdAt: new Date().toISOString(),
		});

		const rootCid = await getRepoRootCid();

		const records = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 10,
		});
		const rkey = records.records[0]!.uri.split("/").pop()!;
		const recordPath = `app.bsky.feed.post/${rkey}`;

		const proof = await generateMstProof(
			ipfsService,
			rootCid,
			recordPath,
		);

		// Verify with a wrong commit CID
		const fakeCommitCid =
			"bafyreig6mxqmjlb7yjbhhhz6vqmtiw4kgipvhqoowdkggjlpzpd5tcm4";
		const verification = await verifyMstProof(
			proof,
			fakeCommitCid,
			recordPath,
		);

		expect(verification.valid).toBe(false);
		expect(verification.error).toContain("Commit block CID mismatch");
	});

	it("verification fails with tampered node bytes", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Tamper test",
			createdAt: new Date().toISOString(),
		});

		const rootCid = await getRepoRootCid();

		const records = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 10,
		});
		const rkey = records.records[0]!.uri.split("/").pop()!;
		const recordPath = `app.bsky.feed.post/${rkey}`;

		const proof = await generateMstProof(
			ipfsService,
			rootCid,
			recordPath,
		);

		// Tamper with the first MST node
		const tamperedProof = {
			...proof,
			nodes: proof.nodes.map((node, i) => {
				if (i === 0) {
					// Flip a byte
					const tampered = new Uint8Array(node.bytes);
					tampered[tampered.length - 1] =
						(tampered[tampered.length - 1]! ^ 0xff) & 0xff;
					return { ...node, bytes: tampered };
				}
				return node;
			}),
		};

		const verification = await verifyMstProof(
			tamperedProof,
			rootCid,
			recordPath,
		);

		expect(verification.valid).toBe(false);
	});

	it("verification fails with wrong record path", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Wrong path test",
			createdAt: new Date().toISOString(),
		});

		const rootCid = await getRepoRootCid();

		const records = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 10,
		});
		const rkey = records.records[0]!.uri.split("/").pop()!;
		const recordPath = `app.bsky.feed.post/${rkey}`;

		const proof = await generateMstProof(
			ipfsService,
			rootCid,
			recordPath,
		);

		// Verify against a different path — the proof says found=true but
		// the verification will walk the nodes with the wrong path
		const wrongPath = "app.bsky.feed.post/totally-wrong-rkey";
		const verification = await verifyMstProof(
			proof,
			rootCid,
			wrongPath,
		);

		// The proof was generated for a different path, so either:
		// - The verifier will not find the record at the wrong path (found mismatch)
		// - Or the node chain won't be valid
		expect(verification.valid).toBe(false);
	});

	it("verification fails with empty nodes array", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Empty nodes test",
			createdAt: new Date().toISOString(),
		});

		const rootCid = await getRepoRootCid();

		const records = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 10,
		});
		const rkey = records.records[0]!.uri.split("/").pop()!;
		const recordPath = `app.bsky.feed.post/${rkey}`;

		const proof = await generateMstProof(
			ipfsService,
			rootCid,
			recordPath,
		);

		const emptyProof = { ...proof, nodes: [] };
		const verification = await verifyMstProof(
			emptyProof,
			rootCid,
			recordPath,
		);

		expect(verification.valid).toBe(false);
		expect(verification.error).toContain("no MST nodes");
	});

	// ============================================
	// Edge cases
	// ============================================

	it("handles a repo with a single record", async () => {
		await repoManager.putRecord("app.bsky.actor.profile", "self", {
			$type: "app.bsky.actor.profile",
			displayName: "Solo",
		});

		const rootCid = await getRepoRootCid();
		const recordPath = "app.bsky.actor.profile/self";

		const proof = await generateMstProof(
			ipfsService,
			rootCid,
			recordPath,
		);

		expect(proof.found).toBe(true);
		expect(proof.nodes.length).toBeGreaterThanOrEqual(1);

		const verification = await verifyMstProof(proof, rootCid, recordPath);
		expect(verification.valid).toBe(true);
		expect(verification.found).toBe(true);
	});

	it("handles many records to create a deeper tree", async () => {
		// Create 50 records to ensure multi-level MST
		for (let i = 0; i < 50; i++) {
			await repoManager.createRecord("app.bsky.feed.post", undefined, {
				$type: "app.bsky.feed.post",
				text: `Deep tree post ${i}`,
				createdAt: new Date().toISOString(),
			});
		}

		const rootCid = await getRepoRootCid();

		const records = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 100,
		});

		// Test the first, middle, and last record
		const indicesToTest = [
			0,
			Math.floor(records.records.length / 2),
			records.records.length - 1,
		];

		for (const idx of indicesToTest) {
			const record = records.records[idx]!;
			const rkey = record.uri.split("/").pop()!;
			const recordPath = `app.bsky.feed.post/${rkey}`;

			const proof = await generateMstProof(
				ipfsService,
				rootCid,
				recordPath,
			);

			expect(proof.found).toBe(true);
			expect(proof.recordCid).not.toBeNull();

			const verification = await verifyMstProof(
				proof,
				rootCid,
				recordPath,
			);

			expect(verification.valid).toBe(true);
			expect(verification.found).toBe(true);
			expect(verification.recordCid).toBe(proof.recordCid);
		}
	});

	it("generate throws when commit block is missing", async () => {
		await expect(
			generateMstProof(
				ipfsService,
				"bafyreig6mxqmjlb7yjbhhhz6vqmtiw4kgipvhqoowdkggjlpzpd5tcm4",
				"app.bsky.feed.post/abc",
			),
		).rejects.toThrow("Commit block not found");
	});

	it("proof roundtrip: generate then verify maintains consistency", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Roundtrip test",
			createdAt: new Date().toISOString(),
		});

		const rootCid = await getRepoRootCid();

		const records = await repoManager.listRecords("app.bsky.feed.post", {
			limit: 10,
		});
		const rkey = records.records[0]!.uri.split("/").pop()!;
		const existingPath = `app.bsky.feed.post/${rkey}`;
		const missingPath = "app.bsky.feed.post/zzz-does-not-exist";

		// Existence proof roundtrip
		const existProof = await generateMstProof(
			ipfsService,
			rootCid,
			existingPath,
		);
		const existVerify = await verifyMstProof(
			existProof,
			rootCid,
			existingPath,
		);
		expect(existVerify.valid).toBe(true);
		expect(existVerify.found).toBe(true);

		// Non-existence proof roundtrip
		const missingProof = await generateMstProof(
			ipfsService,
			rootCid,
			missingPath,
		);
		const missingVerify = await verifyMstProof(
			missingProof,
			rootCid,
			missingPath,
		);
		expect(missingVerify.valid).toBe(true);
		expect(missingVerify.found).toBe(false);
	});

	// ============================================
	// extractAllRecordPaths
	// ============================================

	it("extractAllRecordPaths returns all record paths", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Post one",
			createdAt: new Date().toISOString(),
		});
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Post two",
			createdAt: new Date().toISOString(),
		});
		await repoManager.putRecord("app.bsky.actor.profile", "self", {
			$type: "app.bsky.actor.profile",
			displayName: "Test User",
		});

		const rootCid = await getRepoRootCid();

		const paths = await extractAllRecordPaths(ipfsService, rootCid);

		// Should find all three records
		expect(paths.length).toBe(3);
		expect(paths.filter((p) => p.startsWith("app.bsky.feed.post/")).length).toBe(2);
		expect(paths).toContain("app.bsky.actor.profile/self");
	});

	it("extractAllRecordPaths returns sorted paths for a large repo", async () => {
		for (let i = 0; i < 30; i++) {
			await repoManager.createRecord("app.bsky.feed.post", undefined, {
				$type: "app.bsky.feed.post",
				text: `Post ${i}`,
				createdAt: new Date().toISOString(),
			});
		}

		const rootCid = await getRepoRootCid();
		const paths = await extractAllRecordPaths(ipfsService, rootCid);

		expect(paths.length).toBe(30);
		// MST in-order walk produces sorted keys
		const sorted = [...paths].sort();
		expect(paths).toEqual(sorted);
	});

	it("extractAllRecordPaths returns empty for missing commit", async () => {
		const paths = await extractAllRecordPaths(
			ipfsService,
			"bafyreig6mxqmjlb7yjbhhhz6vqmtiw4kgipvhqoowdkggjlpzpd5tcm4",
		);
		expect(paths).toEqual([]);
	});

	// ============================================
	// extractAllCids
	// ============================================

	it("extractAllCids returns all reachable CIDs", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "CID walk test",
			createdAt: new Date().toISOString(),
		});
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "Another post",
			createdAt: new Date().toISOString(),
		});

		const rootCid = await getRepoRootCid();

		const cids = await extractAllCids(ipfsService, rootCid);

		// Should include at least: commit CID, MST root, MST nodes, record value CIDs
		expect(cids.size).toBeGreaterThanOrEqual(4);
		// The commit CID itself should be in the set
		expect(cids.has(rootCid)).toBe(true);
	});

	it("extractAllCids grows with more records", async () => {
		await repoManager.createRecord("app.bsky.feed.post", undefined, {
			$type: "app.bsky.feed.post",
			text: "First post",
			createdAt: new Date().toISOString(),
		});

		const rootCid1 = await getRepoRootCid();
		const cids1 = await extractAllCids(ipfsService, rootCid1);

		// Add more records
		for (let i = 0; i < 10; i++) {
			await repoManager.createRecord("app.bsky.feed.post", undefined, {
				$type: "app.bsky.feed.post",
				text: `Additional post ${i}`,
				createdAt: new Date().toISOString(),
			});
		}

		const rootCid2 = await getRepoRootCid();
		const cids2 = await extractAllCids(ipfsService, rootCid2);

		// More records = more CIDs
		expect(cids2.size).toBeGreaterThan(cids1.size);
	});

	it("extractAllCids returns only commit CID for missing data", async () => {
		const fakeCid = "bafyreig6mxqmjlb7yjbhhhz6vqmtiw4kgipvhqoowdkggjlpzpd5tcm4";
		const cids = await extractAllCids(ipfsService, fakeCid);
		// Only the commit CID itself (data can't be loaded)
		expect(cids.size).toBe(1);
		expect(cids.has(fakeCid)).toBe(true);
	});
});