/**
 * MST Path Proof: Generate and verify compact Merkle Search Tree proofs.
 *
 * An MST path proof demonstrates that a specific record exists (or does not
 * exist) in an atproto repo by providing only the MST nodes along the path
 * from the root to the leaf, plus the commit block that binds the MST root
 * to a signed commit CID.
 *
 * Proof structure:
 * - commitBlock: the CBOR-encoded commit object (contains `data` field = MST root CID)
 * - nodes: ordered list of { cid, bytes } for each MST node on the path from root to leaf
 * - recordCid: the CID of the record value if found, or null for non-existence proofs
 * - found: whether the record was found in the MST
 */

import { createHash } from "node:crypto";
import { CID } from "multiformats";
import { cborDecode } from "@atproto/common";
import type { BlockStore } from "../ipfs.js";

// ---------- Types ----------

/** A single block in the proof: its CID and raw CBOR bytes. */
export interface ProofBlock {
	cid: string;
	bytes: Uint8Array;
}

/** A compact MST path proof for a single record path. */
export interface MstProof {
	/** The CBOR-encoded commit block. */
	commitBlock: ProofBlock;
	/** Ordered MST node blocks from root to the deepest node on the path. */
	nodes: ProofBlock[];
	/** CID of the record value if found, null if record does not exist. */
	recordCid: string | null;
	/** Whether the record was found in the MST. */
	found: boolean;
}

/** Decoded MST node data (matches @atproto/repo's NodeData shape). */
interface NodeData {
	l: CID | null;
	e: Array<{
		p: number;
		k: Uint8Array;
		v: CID;
		t: CID | null;
	}>;
}

/** Result of verifying an MST proof. */
export interface MstProofVerification {
	/** Whether the proof is valid. */
	valid: boolean;
	/** The record CID if the proof shows existence. */
	recordCid: string | null;
	/** Whether the record was found (matches proof.found). */
	found: boolean;
	/** Error message if verification failed. */
	error?: string;
}

// ---------- Helpers ----------

/** dag-cbor multicodec code */
const DAG_CBOR_CODE = 0x71;

/** sha2-256 multicodec code */
const SHA2_256_CODE = 0x12;

/**
 * Compute the CID for raw CBOR bytes using dag-cbor codec + sha-256.
 * Uses Node.js built-in crypto for the hash.
 */
function cidForBytes(bytes: Uint8Array): CID {
	const hash = createHash("sha256").update(bytes).digest();

	// Build a multihash: varint(code) + varint(length) + digest
	// sha2-256 code = 0x12, digest length = 32
	const multihash = new Uint8Array(2 + hash.length);
	multihash[0] = SHA2_256_CODE;
	multihash[1] = hash.length;
	multihash.set(hash, 2);

	// Create a CIDv1 with dag-cbor codec
	return CID.create(
		1,
		DAG_CBOR_CODE,
		{ code: SHA2_256_CODE, size: hash.length, digest: hash, bytes: multihash } as Parameters<typeof CID.create>[2],
	);
}

/**
 * Decode CBOR bytes and return the MST NodeData structure.
 */
function decodeNodeData(bytes: Uint8Array): NodeData {
	const raw = cborDecode(bytes) as {
		l: CID | null;
		e: Array<{ p: number; k: Uint8Array; v: CID; t: CID | null }>;
	};
	return raw;
}

/**
 * Reconstruct the full keys from a node's compressed entries.
 * Returns an array of { key, value (record CID), subtree (CID or null) }.
 */
function expandEntries(
	data: NodeData,
): Array<{ key: string; value: CID; subtree: CID | null }> {
	const result: Array<{ key: string; value: CID; subtree: CID | null }> = [];
	let lastKey = "";
	for (const entry of data.e) {
		const keyStr = Buffer.from(entry.k).toString("ascii");
		const key = lastKey.slice(0, entry.p) + keyStr;
		result.push({ key, value: entry.v, subtree: entry.t });
		lastKey = key;
	}
	return result;
}

// ---------- Full MST walk ----------

/**
 * Walk an MST node recursively, collecting all record paths (keys).
 */
async function walkMstNode(
	blockStore: BlockStore,
	nodeCid: CID,
	paths: string[],
): Promise<void> {
	const bytes = await blockStore.getBlock(nodeCid.toString());
	if (!bytes) return;

	const nodeData = decodeNodeData(bytes);
	const entries = expandEntries(nodeData);

	// Visit left subtree first
	if (nodeData.l) {
		await walkMstNode(blockStore, nodeData.l, paths);
	}

	// Visit each entry: collect key, then recurse into right subtree
	for (const entry of entries) {
		paths.push(entry.key);
		if (entry.subtree) {
			await walkMstNode(blockStore, entry.subtree, paths);
		}
	}
}

/**
 * Extract all record paths from a repo by walking the full MST.
 *
 * @param blockStore - Block storage containing the repo blocks
 * @param commitCid - CID of the commit block (repo head)
 * @returns All record paths in the repo (e.g. "app.bsky.feed.post/abc123")
 */
export async function extractAllRecordPaths(
	blockStore: BlockStore,
	commitCid: string,
): Promise<string[]> {
	const commitBytes = await blockStore.getBlock(commitCid);
	if (!commitBytes) return [];

	const commitObj = cborDecode(commitBytes) as { data: CID };
	const mstRootCid = commitObj.data;
	if (!mstRootCid) return [];

	const paths: string[] = [];
	await walkMstNode(blockStore, mstRootCid, paths);
	return paths;
}

/**
 * Walk an MST node recursively, collecting ALL CIDs reachable from it.
 * Collects node CIDs, record value CIDs, and subtree pointer CIDs.
 */
async function walkMstCids(
	blockStore: BlockStore,
	nodeCid: CID,
	cids: Set<string>,
): Promise<void> {
	const cidStr = nodeCid.toString();
	if (cids.has(cidStr)) return; // Already visited
	cids.add(cidStr);

	const bytes = await blockStore.getBlock(cidStr);
	if (!bytes) return;

	const nodeData = decodeNodeData(bytes);

	// Visit left subtree
	if (nodeData.l) {
		await walkMstCids(blockStore, nodeData.l, cids);
	}

	// Visit each entry: collect value CID, recurse into right subtree
	for (const entry of nodeData.e) {
		cids.add(entry.v.toString());
		if (entry.t) {
			await walkMstCids(blockStore, entry.t, cids);
		}
	}
}

/**
 * Extract ALL CIDs referenced by a repo at a given commit.
 *
 * Walks the commit → MST → all reachable CIDs including:
 * - The commit CID itself
 * - The MST root CID
 * - Every MST node CID
 * - Every record value CID
 * - Every subtree pointer CID
 *
 * This produces the complete "live CID set" for a repo. Any CID tracked
 * in replication_blocks but NOT in this set is orphaned from this DID's
 * perspective and can be considered for GC.
 *
 * @param blockStore - Block storage containing the repo blocks
 * @param commitCid - CID of the commit block (repo head)
 * @returns Set of all CIDs referenced by the repo
 */
export async function extractAllCids(
	blockStore: BlockStore,
	commitCid: string,
): Promise<Set<string>> {
	const cids = new Set<string>();
	cids.add(commitCid);

	const commitBytes = await blockStore.getBlock(commitCid);
	if (!commitBytes) return cids;

	const commitObj = cborDecode(commitBytes) as { data: CID };
	const mstRootCid = commitObj.data;
	if (!mstRootCid) return cids;

	await walkMstCids(blockStore, mstRootCid, cids);
	return cids;
}

// ---------- Generation ----------

/**
 * Generate a compact MST path proof for a record path.
 *
 * @param blockStore - Block storage containing the repo blocks
 * @param commitCid - CID of the commit block (repo head)
 * @param recordPath - Record path in the form "collection/rkey"
 * @returns An MstProof containing just the blocks needed to verify the path
 */
export async function generateMstProof(
	blockStore: BlockStore,
	commitCid: string,
	recordPath: string,
): Promise<MstProof> {
	// 1. Fetch and decode the commit block
	const commitBytes = await blockStore.getBlock(commitCid);
	if (!commitBytes) {
		throw new Error(`Commit block not found: ${commitCid}`);
	}

	const commitObj = cborDecode(commitBytes) as {
		did: string;
		version: number;
		data: CID;
		rev: string;
		prev: CID | null;
		sig: Uint8Array;
	};

	const mstRootCid = commitObj.data;
	if (!mstRootCid) {
		throw new Error("Commit block has no data field (MST root)");
	}

	const commitBlock: ProofBlock = {
		cid: commitCid,
		bytes: commitBytes,
	};

	// 2. Walk down the MST collecting nodes on the path
	const nodes: ProofBlock[] = [];
	let currentCid: CID = mstRootCid;
	let found = false;
	let recordCid: string | null = null;

	for (;;) {
		const cidStr = currentCid.toString();
		const nodeBytes = await blockStore.getBlock(cidStr);
		if (!nodeBytes) {
			throw new Error(`MST node block not found: ${cidStr}`);
		}

		nodes.push({ cid: cidStr, bytes: nodeBytes });

		const nodeData = decodeNodeData(nodeBytes);
		const entries = expandEntries(nodeData);

		// Find the first entry whose key is >= recordPath
		const index = entries.findIndex((e) => e.key >= recordPath);

		if (index >= 0 && entries[index]!.key === recordPath) {
			// Found the record at this level
			found = true;
			recordCid = entries[index]!.value.toString();
			break;
		}

		// Determine which subtree to descend into:
		// - index < 0: all keys < recordPath => last entry's right subtree
		// - index === 0: recordPath < first key => left pointer (nodeData.l)
		// - index > 0: recordPath between entries[index-1] and entries[index]
		//   => entries[index-1]'s right subtree
		let nextSubtree: CID | null = null;

		if (index < 0) {
			if (entries.length > 0) {
				nextSubtree = entries[entries.length - 1]!.subtree;
			} else {
				nextSubtree = nodeData.l;
			}
		} else if (index === 0) {
			nextSubtree = nodeData.l;
		} else {
			nextSubtree = entries[index - 1]!.subtree;
		}

		if (nextSubtree) {
			currentCid = nextSubtree;
		} else {
			// No subtree to descend into — record does not exist
			break;
		}
	}

	return {
		commitBlock,
		nodes,
		recordCid,
		found,
	};
}

// ---------- Verification ----------

/**
 * Verify an MST path proof against a commit CID and record path.
 *
 * Verification checks:
 * 1. The commit block's CID matches the expected commitCid
 * 2. The commit's `data` field gives the MST root CID
 * 3. Each MST node's CID matches its content hash
 * 4. The path through the MST nodes correctly leads to the claimed record
 *    (or correctly demonstrates non-existence)
 *
 * @param proof - The MST path proof to verify
 * @param commitCid - Expected commit CID
 * @param recordPath - Record path in the form "collection/rkey"
 * @returns Verification result
 */
export async function verifyMstProof(
	proof: MstProof,
	commitCid: string,
	recordPath: string,
): Promise<MstProofVerification> {
	try {
		// 1. Verify the commit block CID
		const actualCommitCid = cidForBytes(proof.commitBlock.bytes);
		if (actualCommitCid.toString() !== commitCid) {
			return {
				valid: false,
				recordCid: null,
				found: false,
				error: `Commit block CID mismatch: expected ${commitCid}, got ${actualCommitCid.toString()}`,
			};
		}

		// 2. Decode the commit and extract the MST root CID
		const commitObj = cborDecode(proof.commitBlock.bytes) as {
			did: string;
			version: number;
			data: CID;
			rev: string;
		};

		const mstRootCid = commitObj.data;
		if (!mstRootCid) {
			return {
				valid: false,
				recordCid: null,
				found: false,
				error: "Commit block has no data field (MST root)",
			};
		}

		// 3. Verify each node block's CID and walk the path
		if (proof.nodes.length === 0) {
			return {
				valid: false,
				recordCid: null,
				found: false,
				error: "Proof contains no MST nodes",
			};
		}

		// First node must be the MST root
		const firstNodeCid = cidForBytes(proof.nodes[0]!.bytes);
		if (firstNodeCid.toString() !== mstRootCid.toString()) {
			return {
				valid: false,
				recordCid: null,
				found: false,
				error: `First proof node CID (${firstNodeCid.toString()}) does not match MST root (${mstRootCid.toString()})`,
			};
		}

		// Walk through nodes verifying the chain
		let expectedCid = mstRootCid.toString();
		let found = false;
		let recordCid: string | null = null;

		for (let i = 0; i < proof.nodes.length; i++) {
			const node = proof.nodes[i]!;

			// Verify block CID matches content
			const actualCid = cidForBytes(node.bytes);
			if (actualCid.toString() !== expectedCid) {
				return {
					valid: false,
					recordCid: null,
					found: false,
					error: `Node ${i} CID mismatch: expected ${expectedCid}, got ${actualCid.toString()}`,
				};
			}

			const nodeData = decodeNodeData(node.bytes);
			const entries = expandEntries(nodeData);

			// Find record in this node
			const entryIndex = entries.findIndex((e) => e.key >= recordPath);

			if (entryIndex >= 0 && entries[entryIndex]!.key === recordPath) {
				// Found the record
				found = true;
				recordCid = entries[entryIndex]!.value.toString();

				// This should be the last node in the proof
				if (i !== proof.nodes.length - 1) {
					return {
						valid: false,
						recordCid: null,
						found: false,
						error: `Record found at node ${i} but proof has ${proof.nodes.length} nodes`,
					};
				}
				break;
			}

			// Determine next subtree
			let nextSubtree: CID | null = null;
			if (entryIndex < 0) {
				if (entries.length > 0) {
					nextSubtree = entries[entries.length - 1]!.subtree;
				} else {
					nextSubtree = nodeData.l;
				}
			} else if (entryIndex === 0) {
				nextSubtree = nodeData.l;
			} else {
				nextSubtree = entries[entryIndex - 1]!.subtree;
			}

			if (i < proof.nodes.length - 1) {
				// There are more nodes — verify the next one connects
				if (!nextSubtree) {
					return {
						valid: false,
						recordCid: null,
						found: false,
						error: `Node ${i} has no subtree for the path, but proof has more nodes`,
					};
				}
				expectedCid = nextSubtree.toString();
			} else {
				// Last node and record not found — this is a non-existence proof
				// Verify there's no subtree to descend into
				if (nextSubtree) {
					return {
						valid: false,
						recordCid: null,
						found: false,
						error: `Last proof node has subtree for the path — proof is incomplete`,
					};
				}
			}
		}

		// Verify the proof's claims match what we verified
		if (found !== proof.found) {
			return {
				valid: false,
				recordCid: null,
				found: false,
				error: `Proof claims found=${proof.found} but verification found found=${found}`,
			};
		}

		if (found && recordCid !== proof.recordCid) {
			return {
				valid: false,
				recordCid: null,
				found: false,
				error: `Proof claims recordCid=${proof.recordCid} but verification found ${recordCid}`,
			};
		}

		return {
			valid: true,
			recordCid: found ? recordCid : null,
			found,
		};
	} catch (err) {
		return {
			valid: false,
			recordCid: null,
			found: false,
			error: err instanceof Error ? err.message : String(err),
		};
	}
}