import { upgradeWebSocket } from 'hono/bun';
import { context as otelContext, ROOT_CONTEXT } from '@opentelemetry/api';
import { tagRoute } from './_route-meta';
/**
 * Context key for WebSocket close promise.
 * Used by middleware to defer session finalization until WebSocket closes.
 */
export const WS_DONE_PROMISE_KEY = '_wsDonePromise';
/**
 * Creates a WebSocket middleware for handling WebSocket connections.
 *
 * The handler must be **synchronous** — it runs inside Hono's `upgradeWebSocket`
 * factory which must return event handlers synchronously for the HTTP upgrade to
 * complete. Async work should go inside `onOpen`, `onMessage`, or `onClose`
 * callbacks, which are properly awaited by the runtime.
 *
 * Use with router.get() to create a WebSocket endpoint:
 *
 * @example
 * ```typescript
 * // Basic synchronous usage
 * import { createRouter, websocket } from '@agentuity/runtime';
 *
 * const router = createRouter();
 *
 * router.get('/ws', websocket((c, ws) => {
 *   ws.onOpen(() => {
 *     c.var.logger.info('WebSocket opened');
 *     ws.send('Welcome!');
 *   });
 *
 *   ws.onMessage((event) => {
 *     c.var.logger.info('Received:', event.data);
 *     ws.send('Echo: ' + event.data);
 *   });
 *
 *   ws.onClose(() => {
 *     c.var.logger.info('WebSocket closed');
 *   });
 * }));
 * ```
 *
 * @example
 * ```typescript
 * // Async work inside callbacks (correct pattern)
 * router.get('/ws', websocket((c, ws) => {
 *   ws.onOpen(async () => {
 *     const user = await fetchUser(c.var.auth);
 *     ws.send(JSON.stringify({ welcome: user.name }));
 *   });
 *
 *   ws.onMessage(async (event) => {
 *     const result = await processMessage(event.data);
 *     ws.send(JSON.stringify(result));
 *   });
 * }));
 * ```
 *
 * @param handler - Synchronous handler function receiving context and WebSocket connection
 * @returns Hono middleware handler for WebSocket upgrade
 */
export function websocket(handler) {
    const wsHandler = upgradeWebSocket((c) => {
        let openHandler;
        let messageHandler;
        let closeHandler;
        let initialized = false;
        // Create done promise for session lifecycle deferral, but ONLY for actual
        // WebSocket upgrade requests. The factory runs unconditionally for every
        // request hitting this route (Hono calls createEvents before attempting
        // server.upgrade). For non-upgrade HTTP requests, setting the promise would
        // cause the middleware to hang forever waiting for an onClose that never fires.
        let resolveDone;
        const isUpgrade = c.req.header('upgrade')?.toLowerCase() === 'websocket';
        if (isUpgrade) {
            const donePromise = new Promise((resolve) => {
                resolveDone = resolve;
            });
            // Defensive: guard against future code adding rejection paths
            donePromise.catch(() => { });
            // Set on context so middleware defers session finalization until WS closes
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            c.set(WS_DONE_PROMISE_KEY, donePromise);
        }
        const wsConnection = {
            onOpen: (h) => {
                openHandler = h;
            },
            onMessage: (h) => {
                messageHandler = h;
            },
            onClose: (h) => {
                closeHandler = h;
            },
            send: (_data) => {
                // This will be bound to the actual ws in the handlers
            },
        };
        // IMPORTANT: We run in ROOT_CONTEXT (no active OTEL span) to avoid a Bun bug
        // where OTEL-instrumented fetch conflicts with streaming responses.
        // See: https://github.com/agentuity/sdk/issues/471
        // See: https://github.com/oven-sh/bun/issues/24766
        const runHandler = () => {
            otelContext.with(ROOT_CONTEXT, () => {
                handler(c, wsConnection);
            });
            initialized = true;
        };
        runHandler();
        return {
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            onOpen: async (event, ws) => {
                try {
                    wsConnection.send = (data) => ws.send(data);
                    if (openHandler) {
                        await otelContext.with(ROOT_CONTEXT, () => openHandler(event));
                    }
                }
                catch (err) {
                    c.var.logger?.error('WebSocket onOpen error:', err);
                    throw err;
                }
            },
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            onMessage: async (event, ws) => {
                try {
                    if (!initialized) {
                        wsConnection.send = (data) => ws.send(data);
                        runHandler();
                    }
                    if (messageHandler) {
                        await otelContext.with(ROOT_CONTEXT, () => messageHandler(event));
                    }
                }
                catch (err) {
                    c.var.logger?.error('WebSocket onMessage error:', err);
                    throw err;
                }
            },
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            onClose: async (event, _ws) => {
                try {
                    if (closeHandler) {
                        await otelContext.with(ROOT_CONTEXT, () => closeHandler(event));
                    }
                }
                catch (err) {
                    c.var.logger?.error('WebSocket onClose error:', err);
                }
                finally {
                    // Resolve the done promise to trigger session finalization
                    // This must fire even if the user's onClose handler throws
                    resolveDone?.();
                }
            },
        };
    });
    const middleware = (c, next) => wsHandler(c, next);
    return tagRoute(middleware, { type: 'websocket' });
}
//# sourceMappingURL=websocket.js.map