import { safeJsonStringify } from "../../util/log-sanitizer.js";
import { getErrorMeta, getLogger } from "../../util/logger.js";
import { LLMObsTracer, } from "../llmobs/index.js";
import { FSM, } from "./fsm.js";
/**
 * FSM wrapper that emits tracing spans for each state transition.
 *
 * ## Span Hierarchy
 *
 * The TracedFSM creates a hierarchical span structure:
 * - **Root Span** (optional): Tracks the entire FSM lifecycle from initialization to terminal state
 * - **State Spans**: Individual spans for each non-terminal state
 *
 * ## Terminal States
 *
 * States marked as "terminal" in the configuration:
 * - Do not create their own spans
 * - Trigger the completion of the root span (if configured)
 * - Represent final states in the FSM lifecycle
 *
 * ## Error Handling
 *
 * Tracing failures are logged but do not interrupt FSM operation. The FSM will continue
 * to function normally even if the tracing infrastructure is unavailable.
 */
export class TracedFSM extends FSM {
    spans;
    sessionId;
    tracer;
    logger;
    getTags;
    constructor(config, tracer = LLMObsTracer.instance, logger = getLogger()) {
        const { sessionId, createTransitionHandler, createTagProvider, ...rest } = config;
        const wrappedConfig = {
            ...rest,
            createTransitionHandler: (fsm) => {
                const tracedFsm = fsm;
                const originalHandler = createTransitionHandler(tracedFsm);
                return async (params) => {
                    return await tracedFsm.handleTransitionTracing(originalHandler, params);
                };
            },
        };
        super(wrappedConfig);
        this.spans = {
            stateSpan: undefined,
            rootSpan: undefined,
            exportedRootSpan: undefined,
        };
        this.sessionId = sessionId;
        this.tracer = tracer;
        this.logger = logger;
        this.getTags = createTagProvider?.(this) ?? (() => ({}));
    }
    /**
     * Exported representation of the active root span (FSM lifecycle), if any.
     */
    get rootSpan() {
        return this.spans.exportedRootSpan;
    }
    async handleTransitionTracing(next, params) {
        this.finishStateSpanIfOpen(params);
        const options = this.getStateSpanOptions(params.to);
        if (options == "terminal") {
            this.finishRootSpanIfOpen(params);
        }
        else {
            this.startRootSpanIfNeeded(params);
            this.startStateSpan(options, params);
        }
        try {
            await next(params);
        }
        catch (error) {
            this.logger.error("Traced FSM transition threw an error", getErrorMeta(error));
            this.finishStateSpanIfOpen(error);
            this.finishRootSpanIfOpen(error);
            throw error;
        }
    }
    startRootSpanIfNeeded(params) {
        if (this.spans.rootSpan) {
            // Root span already active
            return;
        }
        const config = this.getRootSpanOptions(params.event);
        const span = this.startSpan(config, params);
        this.spans.rootSpan = span;
        this.spans.exportedRootSpan = this.tracer.exportSpan(span);
        this.tracer.rootSpan = span;
    }
    finishRootSpanIfOpen(paramsOrError) {
        const span = this.spans.rootSpan;
        if (!span) {
            // No active root span to finish
            return;
        }
        this.finishSpan(span, paramsOrError);
        this.spans.rootSpan = undefined;
        this.tracer.rootSpan = undefined;
    }
    startStateSpan(options, params) {
        this.spans.stateSpan = this.startSpan(options, params);
    }
    finishStateSpanIfOpen(paramsOrError) {
        const stateSpan = this.spans.stateSpan;
        if (!stateSpan) {
            // No active state span to finish. This should only occur on initial transition.
            return;
        }
        this.finishSpan(stateSpan, paramsOrError);
        this.spans.stateSpan = undefined;
    }
    startSpan(options, params) {
        if (this.sessionId) {
            options.sessionId = this.sessionId;
        }
        const span = this.tracer.startActiveSpan(options);
        const sanitizedEvent = this.sanitizeEvent(params.event);
        this.tracer.annotate(span, {
            inputData: safeJsonStringify(sanitizedEvent),
            tags: { "fsm.fromState": params.from },
        });
        return span;
    }
    finishSpan(span, paramsOrError) {
        const tags = this.getTags();
        if (this.isTransitionParams(paramsOrError)) {
            const sanitizedEvent = this.sanitizeEvent(paramsOrError.event);
            this.tracer.annotate(span, {
                outputData: safeJsonStringify(sanitizedEvent),
                tags: { "fsm.toState": paramsOrError.to, ...tags },
                error: this.isErrorTransitionParams(paramsOrError)
                    ? safeJsonStringify(this.getErrorData(paramsOrError))
                    : undefined,
            });
        }
        else {
            this.tracer.annotate(span, {
                error: paramsOrError,
                tags: { ...tags },
            });
        }
        span.finish();
    }
    isTransitionParams(value) {
        return (value !== null &&
            typeof value === "object" &&
            "event" in value &&
            "from" in value &&
            "to" in value);
    }
    isErrorTransitionParams(value) {
        if (!this.isTransitionParams(value)) {
            return false;
        }
        return "error" in value.event;
    }
    getErrorData(value) {
        if (!this.isErrorTransitionParams(value)) {
            return undefined;
        }
        if ("error" in value.event) {
            return value.event.error;
        }
        return value.event;
    }
}
//# sourceMappingURL=traced-fsm.js.map