import { ClientResponseError, ok } from "@atcute/client";
import { mainSchema as applyWritesSchema } from "@atcute/atproto/types/repo/applyWrites";
import { mainSchema as listRecordsSchema } from "@atcute/atproto/types/repo/listRecords";
import { mainSchema as uploadBlobSchema } from "@atcute/atproto/types/repo/uploadBlob";
import type { OAuthUserAgent } from "@atcute/oauth-browser-client";
import { clear } from "#lib/api/cache.ts";
import { createBobbinClient } from "#lib/api/client.ts";
import { jsonPost } from "#lib/api/_request.ts";
import { nsidOf, type Nsid, type Rkey } from "#lib/api/syntax.ts";
import { createClient, mintServiceAuth, serviceDidForHost } from "#lib/auth/agent.ts";
import { missingPermissions, type Permission } from "#lib/auth/scopes.ts";
import {
	fetchLiarBlob,
	fetchLiarRepo,
	listLiarBlobs,
	readLiarRecords,
	type LiarRecord
} from "#lib/migration/liar.ts";
import { sameOrigin, sleep } from "#lib/migration/status.ts";
import type { Fetch } from "#lib/fetch.ts";

// xrpc-server's default json body limit is 100 kb
const MAX_WRITES = 100;
const MAX_BATCH_BYTES = 90 * 1024;
const LIST_PAGE = 100;
const MIN_RATE_LIMIT_WAIT = 10_000;
const UNMARK = nsidOf("org.tangled.foster.unmarkLiarAccount");

export interface ImportProgress {
	phase: "reading" | "blobs" | "importing" | "finalizing";
	done: number;
	total: number;
	overall: { done: number; total: number };
	// set while the pds holds the import back
	resumesAt?: number;
}

// a record whose blob is gone from every pds still moves, without that blob
export interface LostBlob {
	collection: Nsid;
	rkey: Rkey;
	field: string;
}

export interface ImportResult {
	written: number;
	skipped: number;
	lostBlobs: LostBlob[];
}

export class MigrationScopeError extends Error {
	readonly missing: readonly Permission[];
	constructor(missing: readonly Permission[]) {
		super("Tangled needs your permission to write these records to your PDS.");
		this.name = "MigrationScopeError";
		this.missing = missing;
	}
}

interface Write {
	$type: "com.atproto.repo.applyWrites#create" | "com.atproto.repo.applyWrites#update";
	collection: Nsid;
	rkey: Rkey;
	value: LiarRecord["value"];
}

export const missingImportScopes = (
	scope: string,
	writes: Iterable<Pick<Write, "$type" | "collection">>,
	bobbinUrl: string
): readonly Permission[] => {
	const actions = new Map<Nsid, Set<"create" | "update">>();
	for (const write of writes) {
		const action = write.$type === "com.atproto.repo.applyWrites#update" ? "update" : "create";
		actions.set(write.collection, (actions.get(write.collection) ?? new Set()).add(action));
	}
	return missingPermissions(scope, [
		...[...actions].map(([collection, held]): Permission => ({
			resource: "repo",
			collection,
			actions: [...held]
		})),
		{ resource: "rpc", lxm: UNMARK, aud: serviceDidForHost(bobbinUrl) }
	]);
};

export const batchWrites = (writes: readonly Write[]): Write[][] => {
	const encoder = new TextEncoder();
	const batches: Write[][] = [];
	let current: Write[] = [];
	let bytes = 0;
	for (const write of writes) {
		const size = encoder.encode(JSON.stringify(write)).length;
		if (
			current.length > 0 &&
			(current.length >= MAX_WRITES || bytes + size > MAX_BATCH_BYTES)
		) {
			batches.push(current);
			current = [];
			bytes = 0;
		}
		current.push(write);
		bytes += size;
	}
	if (current.length > 0) batches.push(current);
	return batches;
};

// a profile skipped at /welcome holds nothing but its $type
const isBare = (value: unknown): boolean =>
	typeof value === "object" && value !== null && Object.keys(value).join() === "$type";

// rkey to whether the record there is bare
const existingRkeys = async (
	agent: OAuthUserAgent,
	collection: Nsid,
	signal?: AbortSignal
): Promise<Map<string, boolean>> => {
	const client = createClient(agent);
	const rkeys = new Map<string, boolean>();
	let cursor: string | undefined;
	do {
		const page = await ok(
			client.call(listRecordsSchema, {
				params: { repo: agent.sub, collection, limit: LIST_PAGE, cursor },
				signal
			})
		);
		for (const record of page.records)
			rkeys.set(record.uri.slice(record.uri.lastIndexOf("/") + 1), isBare(record.value));
		cursor = page.records.length === 0 ? undefined : page.cursor;
	} while (cursor);
	return rkeys;
};

// existing records win so a retry never rewinds an edit
const writesToMake = async (
	agent: OAuthUserAgent,
	records: readonly LiarRecord[],
	signal?: AbortSignal
): Promise<Write[]> => {
	const collections = [...new Set(records.map((record) => record.collection))];
	const existing = new Map<Nsid, Map<string, boolean>>(
		await Promise.all(
			collections.map(
				async (collection) =>
					[collection, await existingRkeys(agent, collection, signal)] as const
			)
		)
	);
	return records.flatMap((record): Write[] => {
		const bare = existing.get(record.collection)?.get(record.rkey);
		if (bare === false) return [];
		return [
			{
				$type:
					bare === undefined
						? "com.atproto.repo.applyWrites#create"
						: "com.atproto.repo.applyWrites#update",
				collection: record.collection,
				rkey: record.rkey,
				value: record.value
			}
		];
	});
};

// the pds sniffs an upload's type and refuses a ref that names another one
const withStoredMimeTypes = (value: unknown, stored: ReadonlyMap<string, string>): unknown => {
	if (Array.isArray(value)) return value.map((item) => withStoredMimeTypes(item, stored));
	if (value === null || typeof value !== "object") return value;
	// cid links and bytes are class instances that serialize themselves
	if (Object.getPrototypeOf(value) !== Object.prototype) return value;
	const object = value as Record<string, unknown>;
	if (object.$type === "blob") {
		const link = (object.ref as { $link?: string } | undefined)?.$link;
		const mimeType = link === undefined ? undefined : stored.get(link);
		return mimeType === undefined ? object : { ...object, mimeType };
	}
	return Object.fromEntries(
		Object.entries(object).map(([key, item]) => [key, withStoredMimeTypes(item, stored)])
	);
};

const blobLink = (value: unknown): string | undefined =>
	typeof value === "object" && value !== null && (value as { $type?: unknown }).$type === "blob"
		? (value as { ref?: { $link?: string } }).ref?.$link
		: undefined;

const isPlain = (value: unknown): value is Record<string, unknown> =>
	typeof value === "object" &&
	value !== null &&
	Object.getPrototypeOf(value) === Object.prototype;

// the pds names one missing blob per refusal
const missingBlob = (cause: unknown): string | undefined =>
	cause instanceof ClientResponseError && cause.error === "BlobNotFound"
		? /blob: (\S+)/.exec(cause.description ?? "")?.[1]
		: undefined;

const blobFields = (write: Write, cid: string): LostBlob[] => {
	const found: LostBlob[] = [];
	const walk = (value: unknown, path: string) => {
		if (blobLink(value) === cid) {
			found.push({ collection: write.collection, rkey: write.rkey, field: path });
		} else if (Array.isArray(value)) {
			value.forEach((item) => walk(item, path));
		} else if (isPlain(value)) {
			for (const [key, item] of Object.entries(value))
				walk(item, path ? `${path}.${key}` : key);
		}
	};
	walk(write.value, "");
	return found;
};

const withoutBlob = (value: unknown, cid: string): unknown => {
	if (Array.isArray(value)) {
		return value.filter((item) => blobLink(item) !== cid).map((item) => withoutBlob(item, cid));
	}
	if (!isPlain(value)) return value;
	return Object.fromEntries(
		Object.entries(value)
			.filter(([, item]) => blobLink(item) !== cid)
			.map(([key, item]) => [key, withoutBlob(item, cid)])
	);
};

// a pds caps writes per hour, so a big import waits out the window instead of failing
const throttled = async <T>(
	call: () => Promise<T>,
	onWait: (resumesAt: number) => void,
	signal?: AbortSignal
): Promise<T> => {
	for (;;) {
		try {
			return await call();
		} catch (cause) {
			if (!(cause instanceof ClientResponseError) || cause.status !== 429) throw cause;
			const reset = Number(cause.headers.get("ratelimit-reset")) * 1000;
			const resumesAt = Math.max(
				Number.isFinite(reset) ? reset : 0,
				Date.now() + MIN_RATE_LIMIT_WAIT
			);
			onWait(resumesAt);
			await sleep(resumesAt - Date.now(), signal);
		}
	}
};

// goes after the import, hydrant drops the user from liar as soon as this lands
export const unmarkLiarAccount = async (
	agent: OAuthUserAgent,
	bobbinUrl: string,
	{ fetch, signal }: { fetch?: Fetch; signal?: AbortSignal } = {}
): Promise<void> => {
	const token = await mintServiceAuth(agent, {
		aud: serviceDidForHost(bobbinUrl),
		lxm: UNMARK,
		signal
	});
	await jsonPost(createBobbinClient({ serviceUrl: bobbinUrl, fetch }), UNMARK, undefined, {
		headers: { authorization: `Bearer ${token}` },
		signal
	});
};

export const importFromLiar = async (
	agent: OAuthUserAgent,
	{
		liarPdsUrl,
		bobbinUrl,
		fetch,
		signal,
		onProgress
	}: {
		liarPdsUrl: string;
		bobbinUrl: string;
		fetch?: Fetch;
		signal?: AbortSignal;
		onProgress?: (progress: ImportProgress) => void;
	}
): Promise<ImportResult> => {
	if (sameOrigin(agent.session.info.aud, liarPdsUrl)) {
		throw new Error(
			"Your account's PDS is the migration source, so there is nowhere to move your data."
		);
	}

	onProgress?.({ phase: "reading", done: 0, total: 0, overall: { done: 0, total: 0 } });
	const records = await readLiarRecords(
		await fetchLiarRepo(liarPdsUrl, agent.sub, { fetch, signal })
	);

	const writes = await writesToMake(agent, records, signal);

	const missing = missingImportScopes(agent.session.token.scope, writes, bobbinUrl);
	if (missing.length > 0) throw new MigrationScopeError(missing);

	const client = createClient(agent);
	const blobs = await listLiarBlobs(liarPdsUrl, agent.sub, { fetch, signal });
	const batches = batchWrites(writes);
	const overall = { done: 0, total: blobs.length + batches.length };
	const report = (
		phase: "blobs" | "importing",
		done: number,
		total: number,
		resumesAt?: number
	) => onProgress?.({ phase, done, total, overall: { ...overall }, resumesAt });

	const stored = new Map<string, string>();
	report("blobs", 0, blobs.length);
	for (const [i, cid] of blobs.entries()) {
		const input = await fetchLiarBlob(liarPdsUrl, agent.sub, cid, { fetch, signal });
		const { blob } = await throttled(
			() => ok(client.call(uploadBlobSchema, { input, signal })),
			(resumesAt) => report("blobs", i, blobs.length, resumesAt),
			signal
		);
		stored.set(blob.ref.$link, blob.mimeType);
		overall.done++;
		report("blobs", i + 1, blobs.length);
	}

	const lostBlobs: LostBlob[] = [];
	let done = 0;
	report("importing", done, writes.length);
	for (const batch of batches) {
		signal?.throwIfAborted();
		let fixed = batch.map((write) => ({
			...write,
			value: withStoredMimeTypes(write.value, stored) as Write["value"]
		}));
		for (;;) {
			try {
				await throttled(
					() =>
						ok(
							client.call(applyWritesSchema, {
								input: { repo: agent.sub, writes: fixed, validate: false },
								signal
							})
						),
					(resumesAt) => report("importing", done, writes.length, resumesAt),
					signal
				);
				break;
			} catch (cause) {
				const cid = missingBlob(cause);
				if (cid === undefined) throw cause;
				const lost = fixed.flatMap((write) => blobFields(write, cid));
				if (lost.length === 0) throw cause;
				lostBlobs.push(...lost);
				fixed = fixed.map((write) => ({
					...write,
					value: withoutBlob(write.value, cid) as Write["value"]
				}));
			}
		}
		done += batch.length;
		overall.done++;
		report("importing", done, writes.length);
	}

	clear();
	return { written: writes.length, skipped: records.length - writes.length, lostBlobs };
};