import { PgSelect } from 'drizzle-orm/pg-core'
import postgres from 'postgres'

/** Allow using a drizzle query with postgres.js' cursor-based AsyncIterator.
 * To be removed once postgres iterator is added to drizzle
 * https://github.com/drizzle-team/drizzle-orm/issues/456
 *
 * @example
 * const streamQuery = createPostgresDrizzleStreamer(postgresClient)
 * const documentsReader = streamQuery(db.select().from(documents).$dynamic(), 2000)
 * for await (const documentsChunk of documentsReader) {}
 */
export function createPostgresDrizzleStreamer(postgresClient: postgres.Sql) {
    return async function* streamQuery<
        Qb extends PgSelect,
        // Infer the shape of a single result object
        ResultItem extends Qb['_']['result'][number],
        // Define the raw result type from the driver (likely Record<string, any>)
        RawResultItem = Record<string, any>
    >(qb: Qb, chunkSize?: number): AsyncGenerator<ResultItem[]> {
        const q = qb.toSQL();
        // Use a more specific type if possible, but RawResultItem provides flexibility
        const chunkReader = postgresClient
            .unsafe<RawResultItem[]>(
                q.sql,
                q.params as Parameters<typeof postgresClient.unsafe>[1],
            )
            .cursor(chunkSize);

        const selectedFields = qb._.selectedFields; // Cache selected fields info

        try {
            for await (const chunk of chunkReader) {
                // Process the chunk efficiently
                yield chunk.map((rawResult) => {
                    const mappedResult: Record<string, unknown> = {};
                    // Iterate over the keys of the raw result object
                    for (const key in rawResult) {
                        // Ensure the key is directly on the object (not from prototype)
                        if (Object.prototype.hasOwnProperty.call(rawResult, key)) {
                            const col = selectedFields[key];
                            if (col === undefined) {
                                // Handle cases where a key from the raw result doesn't match selected fields
                                // This might happen with complex queries or if the driver returns extra info.
                                // Throwing an error maintains consistency with the original code.
                                throw new Error(`Field '${key}' from query result not found in Drizzle's selected fields.`);
                            }
                            // Map the raw value using Drizzle's type converter
                            mappedResult[key] = col.mapFromDriverValue(rawResult[key]);
                        }
                    }
                    // Cast the mapped object to the expected Drizzle result type
                    return mappedResult as ResultItem;
                });
            }
        } catch (error) {
            // Optional: Add logging or specific error handling for stream errors
            console.error("Error during database stream:", error);
            throw error; // Re-throw the error after logging
        }
    }
}