"use strict";
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
    return new (P || (P = Promise))(function (resolve, reject) {
        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
        step((generator = generator.apply(thisArg, _arguments || [])).next());
    });
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.LlamaParallelAPI = void 0;
const utils_1 = require("./utils");
class LlamaParallelAPI {
    constructor(context) {
        this.enabled = false;
        this.pendingRequests = new Map();
        this.context = context;
    }
    /**
     * Enable parallel decoding mode
     * @param config Configuration for parallel mode
     * @returns boolean indicating if successful
     */
    enable(config) {
        return __awaiter(this, void 0, void 0, function* () {
            const defaultConfig = { n_parallel: 2, n_batch: 512 };
            const result = this.context.enableParallelMode(Object.assign(Object.assign({}, defaultConfig), config));
            this.enabled = result;
            return result;
        });
    }
    /**
     * Disable parallel decoding mode
     */
    disable() {
        this.context.disableParallelMode();
        this.enabled = false;
    }
    /**
     * Configure parallel decoding mode (enables if not already enabled)
     * @param config Configuration for parallel mode
     * @returns boolean indicating if successful
     */
    configure(config) {
        return __awaiter(this, void 0, void 0, function* () {
            return this.enable(config);
        });
    }
    /**
     * Queue a completion request for parallel processing
     * @param options Completion options
     * @param onToken Optional callback for each token
     * @returns Object with requestId, promise for result, and stop function
     */
    completion(options, onToken) {
        return __awaiter(this, void 0, void 0, function* () {
            if (!this.enabled) {
                throw new Error('Parallel mode is not enabled. Call enable() first.');
            }
            const tokenCallback = onToken
                ? (error, result) => {
                    if (error) {
                        console.error('Token callback error:', error);
                        // Handle completion error
                        const pendingReq = this.pendingRequests.get(result === null || result === void 0 ? void 0 : result.requestId);
                        if (pendingReq) {
                            pendingReq.reject(error);
                            this.pendingRequests.delete(result === null || result === void 0 ? void 0 : result.requestId);
                        }
                        return;
                    }
                    // Check if this is a token callback or final result
                    if (result) {
                        if (result.token !== undefined) {
                            // This is a token callback
                            onToken(result.requestId, result);
                        }
                        else if (result.text !== undefined ||
                            result.content !== undefined) {
                            // This is the final result
                            const pendingReq = this.pendingRequests.get(result.requestId);
                            if (pendingReq) {
                                pendingReq.resolve(result);
                                this.pendingRequests.delete(result.requestId);
                            }
                        }
                    }
                }
                : undefined;
            const { messages, media_paths = options.media_paths } = (0, utils_1.formatMediaChat)(options.messages);
            // Queue the completion immediately (this is synchronous!)
            const { requestId } = this.context.queueCompletion(Object.assign(Object.assign({}, options), { messages, media_paths: media_paths }), tokenCallback ||
                ((error, result) => {
                    if (error) {
                        const pendingReq = this.pendingRequests.get(result === null || result === void 0 ? void 0 : result.requestId);
                        if (pendingReq) {
                            pendingReq.reject(error);
                            this.pendingRequests.delete(result === null || result === void 0 ? void 0 : result.requestId);
                        }
                    }
                    else if (result &&
                        (result.text !== undefined || result.content !== undefined)) {
                        // Final result for non-streaming
                        const pendingReq = this.pendingRequests.get(result.requestId);
                        if (pendingReq) {
                            pendingReq.resolve(result);
                            this.pendingRequests.delete(result.requestId);
                        }
                    }
                }));
            // Create promise for final result
            const promise = new Promise((resolveResult, rejectResult) => {
                this.pendingRequests.set(requestId, {
                    resolve: resolveResult,
                    reject: rejectResult,
                });
            });
            // Create stop function
            const stop = () => {
                this.context.cancelRequest(requestId);
                const pendingReq = this.pendingRequests.get(requestId);
                if (pendingReq) {
                    pendingReq.reject(new Error('Request cancelled'));
                    this.pendingRequests.delete(requestId);
                }
            };
            // Return immediately without wrapping in a Promise
            return {
                requestId,
                promise,
                stop,
            };
        });
    }
    /**
     * Queue an embedding request for parallel processing
     * @param text Text to embed
     * @param params Optional embedding parameters
     * @returns Object with requestId and promise for result
     */
    embedding(text, params) {
        return __awaiter(this, void 0, void 0, function* () {
            if (!this.enabled) {
                throw new Error('Parallel mode is not enabled. Call enable() first.');
            }
            // Create promise for result
            let resolveResult;
            let rejectResult;
            const promise = new Promise((res, rej) => {
                resolveResult = res;
                rejectResult = rej;
            });
            // Queue the embedding immediately (this is synchronous!)
            const { requestId } = this.context.queueEmbedding(text, params, (error, result) => {
                if (error) {
                    rejectResult(error);
                }
                else {
                    resolveResult(result);
                }
            });
            // Return immediately without wrapping in a Promise
            return {
                requestId,
                promise,
            };
        });
    }
    /**
     * Queue a rerank request for parallel processing
     * @param query Query text
     * @param documents Documents to rank
     * @param params Optional rerank parameters
     * @returns Object with requestId and promise for results
     */
    rerank(query, documents, params) {
        return __awaiter(this, void 0, void 0, function* () {
            if (!this.enabled) {
                throw new Error('Parallel mode is not enabled. Call enable() first.');
            }
            // Create promise for result
            let resolveResult;
            let rejectResult;
            const promise = new Promise((res, rej) => {
                resolveResult = res;
                rejectResult = rej;
            });
            // Queue the rerank immediately (this is synchronous!)
            const { requestId } = this.context.queueRerank(query, documents, params, (error, result) => {
                if (error) {
                    rejectResult(error);
                }
                else {
                    // Add document text to results and sort by score
                    const enrichedResults = result.results
                        .map((r) => (Object.assign(Object.assign({}, r), { document: documents[r.index] })))
                        .sort((a, b) => b.score - a.score);
                    resolveResult(enrichedResults);
                }
            });
            // Return immediately without wrapping in a Promise
            return {
                requestId,
                promise,
            };
        });
    }
    /**
     * Check if parallel mode is enabled
     */
    isEnabled() {
        return this.enabled;
    }
    /**
     * Get current parallel processing status (one-time snapshot)
     * @returns Current parallel status
     */
    getStatus() {
        if (!this.enabled) {
            throw new Error('Parallel mode is not enabled. Call enable() first.');
        }
        return this.context.getParallelStatus();
    }
    /**
     * Subscribe to parallel processing status changes
     * @param callback Called whenever parallel status changes
     * @returns Object with remove() method to unsubscribe
     */
    subscribeToStatus(callback) {
        if (!this.enabled) {
            throw new Error('Parallel mode is not enabled. Call enable() first.');
        }
        const { subscriberId } = this.context.subscribeParallelStatus(callback);
        return {
            remove: () => {
                this.context.unsubscribeParallelStatus(subscriberId);
            },
        };
    }
}
exports.LlamaParallelAPI = LlamaParallelAPI;
