import type { PropsWithChildren } from "react";
import {
	createContext,
	useMemo,
	useCallback,
	useEffect,
	useState,
	useRef,
} from "react";
import {
	createCollection,
	type InferSchemaInput,
	type Collection,
	type InferSchemaOutput,
	type SyncMode,
} from "@tanstack/db";
import { getTableName, type Table } from "drizzle-orm";
import {
	drizzleIndexedDBCollectionOptions,
	type DrizzleIndexedDBCollectionConfig,
} from "../collections/drizzle-indexeddb-collection";
import type { CollectionUtils } from "@firtoz/db-helpers";
import type {
	IdOf,
	InsertSchema,
	SelectSchema,
	GetTableFromSchema,
	InferCollectionFromTable,
} from "@firtoz/drizzle-utils";
import {
	type Migration,
	migrateIndexedDBWithFunctions,
} from "../function-migrator";
import type { IDBCreator, IDBDatabaseLike } from "../idb-types";
import { openIndexedDb } from "../idb-operations";

interface CollectionCacheEntry {
	// biome-ignore lint/suspicious/noExplicitAny: Cache needs to store collections of various types
	collection: Collection<any, string, CollectionUtils<any>, any, any>;
	refCount: number;
}

// Type for migration functions (generated by Drizzle)

export type IndexedDbCollection<
	TSchema extends Record<string, unknown>,
	TTableName extends keyof TSchema & string,
> = Collection<
	InferSchemaOutput<SelectSchema<GetTableFromSchema<TSchema, TTableName>>>,
	IdOf<GetTableFromSchema<TSchema, TTableName>>,
	CollectionUtils<
		InferSchemaOutput<SelectSchema<GetTableFromSchema<TSchema, TTableName>>>
	>,
	SelectSchema<GetTableFromSchema<TSchema, TTableName>>,
	InferSchemaInput<InsertSchema<GetTableFromSchema<TSchema, TTableName>>>
>;

export type DrizzleIndexedDBContextValue<
	TSchema extends Record<string, unknown>,
> = {
	indexedDB: IDBDatabaseLike | null;
	getCollection: <TTableName extends keyof TSchema & string>(
		tableName: TTableName,
	) => IndexedDbCollection<TSchema, TTableName>;
	incrementRefCount: (tableName: string) => void;
	decrementRefCount: (tableName: string) => void;
};

export const DrizzleIndexedDBContext =
	// biome-ignore lint/suspicious/noExplicitAny: Context needs to accept any schema type
	createContext<DrizzleIndexedDBContextValue<any> | null>(null);

type DrizzleIndexedDBProviderProps<TSchema extends Record<string, unknown>> =
	PropsWithChildren<{
		dbName: string;
		schema: TSchema;
		migrations?: Migration[];
		migrateFunction?: (
			dbName: string,
			migrations: Migration[],
			debug?: boolean,
			dbCreator?: IDBCreator,
		) => Promise<IDBDatabaseLike>;
		debug?: boolean;
		syncMode?: SyncMode;
		/**
		 * Optional custom database creator for testing/mocking.
		 * Use createInstrumentedDbCreator() to track IndexedDB operations.
		 */
		dbCreator?: IDBCreator;
	}>;

export function DrizzleIndexedDBProvider<
	TSchema extends Record<string, unknown>,
>({
	children,
	dbName,
	schema,
	migrations = [],
	migrateFunction = migrateIndexedDBWithFunctions,
	debug = false,
	syncMode = "eager",
	dbCreator,
}: DrizzleIndexedDBProviderProps<TSchema>) {
	const [indexedDB, setIndexedDB] = useState<IDBDatabaseLike | null>(null);
	const indexedDBRef = useRef<IDBDatabaseLike | null>(null);
	const [readyPromise] = useState(() => {
		let resolveReady: () => void;
		const promise = new Promise<void>((resolve) => {
			resolveReady = resolve;
		});
		// biome-ignore lint/style/noNonNullAssertion: resolveReady is guaranteed to be set
		return { promise, resolve: resolveReady! };
	});

	useEffect(() => {
		let db: IDBDatabaseLike | null = null;
		const initDB = async () => {
			try {
				if (migrations.length === 0) {
					// Open database directly without migration logic
					db = await openIndexedDb(dbName, dbCreator);
				} else {
					db = await migrateFunction(dbName, migrations, debug, dbCreator);
				}

				indexedDBRef.current = db;
				setIndexedDB(db);
				readyPromise.resolve();
			} catch (error) {
				console.error(
					`[DrizzleIndexedDBProvider] Failed to initialize database ${dbName}:`,
					error,
				);
				throw error;
			}
		};

		initDB();

		// Cleanup on unmount
		return () => {
			if (db) {
				db.close();
			}
		};
	}, [dbName, migrations, migrateFunction, debug, readyPromise, dbCreator]);

	// Collection cache with ref counting
	const collections = useMemo<Map<string, CollectionCacheEntry>>(
		() => new Map(),
		[],
	);

	const getCollection = useCallback<
		DrizzleIndexedDBContextValue<TSchema>["getCollection"]
	>(
		<TTableName extends keyof TSchema & string>(tableName: TTableName) => {
			const cacheKey = tableName;

			// Check if collection already exists in cache
			if (!collections.has(cacheKey)) {
				// Get the table definition from schema
				const table = schema[tableName] as Table;

				if (!table) {
					throw new Error(
						`Table "${tableName}" not found in schema. Available tables: ${Object.keys(schema).join(", ")}`,
					);
				}

				// Extract the actual store/table name from the table definition
				const actualTableName = getTableName(table);

				// Create collection options
				const collectionConfig = drizzleIndexedDBCollectionOptions({
					indexedDBRef,
					table,
					storeName: actualTableName,
					readyPromise: readyPromise.promise,
					debug,
					syncMode,
				} as DrizzleIndexedDBCollectionConfig<Table>);

				// Create new collection and cache it with ref count 0
				// The collection will wait for readyPromise before accessing the database
				// Cast is safe: our collectionConfig provides CollectionUtils, but createCollection types it as UtilsRecord
				const collection = createCollection(
					collectionConfig,
				) as CollectionCacheEntry["collection"];

				collections.set(cacheKey, {
					collection,
					refCount: 0,
				});
			}

			// biome-ignore lint/style/noNonNullAssertion: We just ensured the collection exists
			return collections.get(cacheKey)!
				.collection as unknown as IndexedDbCollection<TSchema, TTableName>;
		},
		[collections, schema, readyPromise.promise, debug, syncMode],
	);

	const incrementRefCount: DrizzleIndexedDBContextValue<TSchema>["incrementRefCount"] =
		useCallback(
			(tableName: string) => {
				const entry = collections.get(tableName);
				if (entry) {
					entry.refCount++;
				}
			},
			[collections],
		);

	const decrementRefCount: DrizzleIndexedDBContextValue<TSchema>["decrementRefCount"] =
		useCallback(
			(tableName: string) => {
				const entry = collections.get(tableName);
				if (entry) {
					entry.refCount--;

					// If ref count reaches 0, remove from cache
					if (entry.refCount <= 0) {
						collections.delete(tableName);
					}
				}
			},
			[collections],
		);

	const contextValue: DrizzleIndexedDBContextValue<TSchema> = useMemo(
		() => ({
			indexedDB,
			getCollection,
			incrementRefCount,
			decrementRefCount,
		}),
		[indexedDB, getCollection, incrementRefCount, decrementRefCount],
	);

	return (
		<DrizzleIndexedDBContext.Provider value={contextValue}>
			{children}
		</DrizzleIndexedDBContext.Provider>
	);
}

// Hook that components use to get a collection with automatic ref counting
export function useIndexedDBCollection<
	TSchema extends Record<string, unknown>,
	TTableName extends keyof TSchema & string,
>(
	context: DrizzleIndexedDBContextValue<TSchema>,
	tableName: TTableName,
): InferCollectionFromTable<GetTableFromSchema<TSchema, TTableName>> {
	const { collection, unsubscribe } = useMemo(() => {
		// Get the collection and increment ref count
		const col = context.getCollection(tableName);
		context.incrementRefCount(tableName);

		// Return collection and unsubscribe function
		return {
			collection: col,
			unsubscribe: () => {
				context.decrementRefCount(tableName);
			},
		};
	}, [context, tableName]);

	// Cleanup on unmount
	useEffect(() => {
		return () => {
			unsubscribe();
		};
	}, [unsubscribe]);

	return collection as unknown as InferCollectionFromTable<
		GetTableFromSchema<TSchema, TTableName>
	>;
}