import type { RunMetadata } from "./runtime-metadata.js";

export type WorkflowStateValue =
	| null
	| boolean
	| number
	| string
	| ReadonlyArray<WorkflowStateValue>
	| { readonly [key: string]: WorkflowStateValue };

export type WorkflowState = {
	get(key: string): Promise<WorkflowStateValue | undefined>;
	set(key: string, value: WorkflowStateValue | undefined): void;
};

type WorkflowStateMutation =
	| { operation: "set"; key: string; value: WorkflowStateValue }
	| { operation: "delete"; key: string };

type WorkflowStepScope = {
	id: string;
	active: boolean;
	mutations: Map<string, WorkflowStateMutation>;
};

const MAX_STATE_KEY_BYTES = 256;
const WORKFLOW_CAPABILITY_KEY_ENV = "NOTION_WORKFLOW_CAPABILITY_KEY";

export function createWorkflowStepState(
	metadata: Pick<RunMetadata, "workerId">,
	stepId: string,
): {
	state: WorkflowState;
	run<T>(fn: () => T | Promise<T>): Promise<T>;
} {
	const scope: WorkflowStepScope = {
		id: stepId,
		active: true,
		mutations: new Map(),
	};
	const state: WorkflowState = {
		async get(key: string): Promise<WorkflowStateValue | undefined> {
			requireActiveStep(scope, "get");
			validateStateKey(key);
			const pending = scope.mutations.get(key);
			if (pending?.operation === "set") {
				return pending.value;
			}
			if (pending?.operation === "delete") {
				return undefined;
			}

			const response = await requestWorkflowState(metadata, {
				operation: "get",
				key,
			});
			const body: { exists: boolean; value?: WorkflowStateValue } = JSON.parse(
				await response.text(),
			);
			if (!body.exists) {
				return undefined;
			}
			return body.value;
		},
		set(key: string, value: WorkflowStateValue | undefined): void {
			requireActiveStep(scope, "set");
			validateStateKey(key);
			if (value === undefined) {
				recordMutation(scope, { operation: "delete", key });
				return;
			}
			recordMutation(scope, { operation: "set", key, value });
		},
	};

	return {
		state,
		async run<T>(fn: () => T | Promise<T>): Promise<T> {
			try {
				const result = await fn();
				if (scope.mutations.size > 0) {
					await requestWorkflowState(metadata, {
						operation: "commit",
						operationId: scope.id,
						mutations: [...scope.mutations.values()],
					});
				}
				return result;
			} finally {
				scope.active = false;
			}
		},
	};
}

function requireActiveStep(expectedScope: WorkflowStepScope, operation: keyof WorkflowState): void {
	if (!expectedScope.active) {
		throw new Error(`state.${operation}() must be called inside context.step().`);
	}
}

function recordMutation(scope: WorkflowStepScope, mutation: WorkflowStateMutation): void {
	scope.mutations.set(mutation.key, mutation);
}

async function requestWorkflowState(
	metadata: Pick<RunMetadata, "workerId">,
	request:
		| { operation: "get"; key: string }
		| {
				operation: "commit";
				operationId: string;
				mutations: WorkflowStateMutation[];
		  },
): Promise<Response> {
	const baseUrl = process.env.NOTION_API_BASE_URL;
	const token = process.env.NOTION_API_TOKEN;
	const capabilityKey = process.env[WORKFLOW_CAPABILITY_KEY_ENV];
	if (!baseUrl || !token || !capabilityKey) {
		throw new Error(
			`Workflow state is only available when the workflow runtime provides NOTION_API_BASE_URL, NOTION_API_TOKEN, and ${WORKFLOW_CAPABILITY_KEY_ENV}.`,
		);
	}

	const response = await fetch(`${baseUrl.replace(/\/$/, "")}/api/v3/workersWorkflowState`, {
		method: "POST",
		headers: {
			authorization: `Bearer ${token}`,
			"content-type": "application/json",
		},
		body: JSON.stringify({
			...request,
			workerId: metadata.workerId,
			capabilityKey,
		}),
	});
	if (!response.ok) {
		throw new Error(
			`Workflow state ${request.operation} failed with status ${response.status}.`,
		);
	}
	return response;
}

function validateStateKey(key: string): void {
	if (key.length === 0) {
		throw new Error("Workflow state keys must not be empty.");
	}
	if (Buffer.byteLength(key) > MAX_STATE_KEY_BYTES) {
		throw new Error(`Workflow state keys must be at most ${MAX_STATE_KEY_BYTES} UTF-8 bytes.`);
	}
}