import SSHClient from '@dylankenneally/react-native-ssh-sftp';
import { queryOptions } from '@tanstack/react-query';
import * as SecureStore from 'expo-secure-store';
import * as z from 'zod';
import { queryClient } from './utils';

const keys = {
	storagePrefix: 'privateKey_',
	manifestKey: 'privateKeysManifest',
	chunkSize: 1800, // Safely under 2048 byte limit
} as const;

const keyManifestSchema = z.object({
	manifestVersion: z.number().default(1),
	keys: z.array(
		z.object({
			id: z.string(),
			priority: z.number(),
			createdAtMs: z.int(),
			chunkCount: z.number().default(1),
		}),
	),
});

async function getKeyManifest() {
	const rawManifest = await SecureStore.getItemAsync(keys.manifestKey);
	const manifest = rawManifest
		? JSON.parse(rawManifest)
		: {
				manifestVersion: 1,
				keys: [],
			};
	return keyManifestSchema.parse(manifest);
}

// Utility functions for chunking large data
function splitIntoChunks(data: string, chunkSize: number): string[] {
	const chunks: string[] = [];
	for (let i = 0; i < data.length; i += chunkSize) {
		chunks.push(data.substring(i, i + chunkSize));
	}
	return chunks;
}

function getChunkKey(keyId: string, chunkIndex: number): string {
	return `${keys.storagePrefix}${keyId}_chunk_${chunkIndex}`;
}

async function savePrivateKey(params: {
	keyId: string;
	privateKey: string;
	priority: number;
}) {
	const manifest = await getKeyManifest();

	const existingKey = manifest.keys.find((key) => key.id === params.keyId);

	if (existingKey) throw new Error('Key already exists');

	// Split the private key into chunks if it's too large
	const chunks = splitIntoChunks(params.privateKey, keys.chunkSize);
	const chunkCount = chunks.length;

	const newKey = {
		id: params.keyId,
		priority: params.priority,
		createdAtMs: Date.now(),
		chunkCount,
	};

	manifest.keys.push(newKey);

	// Save each chunk separately
	for (let i = 0; i < chunks.length; i++) {
		await SecureStore.setItemAsync(getChunkKey(params.keyId, i), chunks[i]!);
	}

	await SecureStore.setItemAsync(keys.manifestKey, JSON.stringify(manifest));
	await queryClient.invalidateQueries({ queryKey: [keyQueryKey] });
}

async function getPrivateKey(keyId: string) {
	const manifest = await getKeyManifest();
	const key = manifest.keys.find((key) => key.id === keyId);
	if (!key) throw new Error('Key not found');

	// Reassemble the private key from chunks
	const chunks: string[] = [];
	for (let i = 0; i < key.chunkCount; i++) {
		const chunk = await SecureStore.getItemAsync(getChunkKey(keyId, i));
		if (!chunk) throw new Error(`Key chunk ${i} not found`);
		chunks.push(chunk);
	}

	const privateKey = chunks.join('');
	return {
		...key,
		privateKey,
	};
}

async function deletePrivateKey(keyId: string) {
	const manifest = await getKeyManifest();
	const key = manifest.keys.find((key) => key.id === keyId);
	if (!key) throw new Error('Key not found');

	manifest.keys = manifest.keys.filter((key) => key.id !== keyId);

	// Delete all chunks for this key
	for (let i = 0; i < key.chunkCount; i++) {
		await SecureStore.deleteItemAsync(getChunkKey(keyId, i));
	}

	await SecureStore.setItemAsync(keys.manifestKey, JSON.stringify(manifest));
	await queryClient.invalidateQueries({ queryKey: [keyQueryKey] });
}

const connections = {
	storagePrefix: 'connection_',
	manifestKey: 'connectionsManifest',
} as const;

const connectionsManifestSchema = z.object({
	manifestVersion: z.number().default(1),
	connections: z.array(
		z.object({
			id: z.string(),
			priority: z.number(),
			createdAt: z.date(),
			modifiedAt: z.date(),
		}),
	),
});

export const connectionDetailsSchema = z.object({
	host: z.string().min(1),
	port: z.number().min(1),
	username: z.string().min(1),
	security: z.discriminatedUnion('type', [
		z.object({
			type: z.literal('password'),
			password: z.string().min(1),
		}),
		z.object({
			type: z.literal('key'),
			keyId: z.string().min(1),
		}),
	]),
});

export type ConnectionDetails = z.infer<typeof connectionDetailsSchema>;

async function getConnectionManifest() {
	const rawManifest = await SecureStore.getItemAsync(connections.manifestKey);
	const manifest = rawManifest
		? JSON.parse(rawManifest)
		: {
				manifestVersion: 1,
				connections: [],
			};
	return connectionsManifestSchema.parse(manifest);
}

async function upsertConnection(params: {
	id: string;
	details: ConnectionDetails;
	priority: number;
}) {
	const manifest = await getConnectionManifest();
	const existingConnection = manifest.connections.find(
		(connection) => connection.id === params.id,
	);

	const newConnection = existingConnection
		? {
				...existingConnection,
				priority: params.priority,
				modifiedAt: new Date(),
			}
		: {
				id: params.id,
				priority: params.priority,
				createdAt: new Date(),
				modifiedAt: new Date(),
			};

	await SecureStore.setItemAsync(
		connections.manifestKey,
		JSON.stringify(manifest),
	);
	await SecureStore.setItemAsync(
		`${connections.storagePrefix}${params.id}`,
		JSON.stringify(params.details),
	);
	await queryClient.invalidateQueries({ queryKey: [connectionQueryKey] });
	return existingConnection ?? newConnection;
}

async function deleteConnection(id: string) {
	const manifest = await getConnectionManifest();
	const connection = manifest.connections.find(
		(connection) => connection.id === id,
	);
	if (!connection) throw new Error('Connection not found');
	manifest.connections = manifest.connections.filter(
		(connection) => connection.id !== id,
	);
	await SecureStore.setItemAsync(
		connections.manifestKey,
		JSON.stringify(manifest),
	);
	await SecureStore.deleteItemAsync(`${connections.storagePrefix}${id}`);
	await queryClient.invalidateQueries({ queryKey: [connectionQueryKey] });
}

async function getConnection(id: string) {
	const manifest = await getConnectionManifest();
	const connection = manifest.connections.find(
		(connection) => connection.id === id,
	);
	if (!connection) throw new Error('Connection not found');
	const detailsString = await SecureStore.getItemAsync(
		`${connections.storagePrefix}${id}`,
	);
	if (!detailsString) throw new Error('Connection details not found');
	const detailsJson = JSON.parse(detailsString);
	const details = connectionDetailsSchema.parse(detailsJson);
	return { ...connection, details };
}

const connectionQueryKey = 'connections';

const listConnectionsQueryOptions = queryOptions({
	queryKey: [connectionQueryKey],
	queryFn: async () => {
		const manifest = await getConnectionManifest();
		const firstConnectionMeta = manifest.connections[0];
		const firstConnection = firstConnectionMeta
			? await getConnection(firstConnectionMeta.id)
			: null;

		return {
			manifest,
			firstConnection,
		};
	},
});

const getConnectionQueryOptions = (id: string) =>
	queryOptions({
		queryKey: [connectionQueryKey, id],
		queryFn: () => getConnection(id),
	});

const keyQueryKey = 'keys';

const listKeysQueryOptions = queryOptions({
	queryKey: [keyQueryKey],
	queryFn: getKeyManifest,
});

// https://github.com/dylankenneally/react-native-ssh-sftp/blob/ea55436d8d40378a8f9dabb95b463739ffb219fa/android/src/main/java/me/keeex/rnssh/RNSshClientModule.java#L101-L119
export type SshPrivateKeyType = 'dsa' | 'rsa' | 'ecdsa' | 'ed25519' | 'ed448';
async function generateKeyPair(params: {
	type: SshPrivateKeyType;
	passphrase?: string;
	keySize?: number;
	comment?: string;
}) {
	const keyPair = await SSHClient.generateKeyPair(
		params.type,
		params.passphrase ?? '',
		params.keySize,
		params.comment ?? '',
	);
	return keyPair;
}

export const secretsManager = {
	keys: {
		utils: {
			getKeyManifest,
			savePrivateKey,
			getPrivateKey,
			deletePrivateKey,
			generateKeyPair,
		},
		query: {
			list: listKeysQueryOptions,
		},
	},
	connections: {
		utils: {
			upsertConnection,
			deleteConnection,
		},
		query: {
			list: listConnectionsQueryOptions,
			get: getConnectionQueryOptions,
		},
	},
};
