// lib/v3/handlers/extractHandler.ts
import { extract as runExtract } from "../../inference.js";
import { getZFactory, getZodType, injectUrls, transformSchema, } from "../../utils.js";
import { v3Logger } from "../logger.js";
import { V3FunctionName } from "../types/public/methods.js";
import { captureHybridSnapshot } from "../understudy/a11y/snapshot/index.js";
import { defaultExtractSchema, pageTextSchema, } from "../types/public/methods.js";
import { StagehandInvalidArgumentError, ExtractTimeoutError, } from "../types/public/sdkErrors.js";
import { createTimeoutGuard } from "./handlerUtils/timeoutGuard.js";
/**
 * Scans the provided Zod schema for any `z.string().url()` fields and
 * replaces them with `z.number()`.
 *
 * @param schema - The Zod object schema to transform.
 * @returns A tuple containing:
 *   1. The transformed schema (or the original schema if no changes were needed).
 *   2. An array of {@link ZodPathSegments} objects representing all the replaced URL fields,
 *      with each path segment showing where in the schema the replacement occurred.
 */
export function transformUrlStringsToNumericIds(schema) {
    const [finalSchema, urlPaths] = transformSchema(schema, []);
    return [finalSchema, urlPaths];
}
export class ExtractHandler {
    llmClient;
    defaultModelName;
    defaultClientOptions;
    resolveLlmClient;
    systemPrompt;
    logInferenceToFile;
    experimental;
    onMetrics;
    constructor(llmClient, defaultModelName, defaultClientOptions, resolveLlmClient, systemPrompt, logInferenceToFile, experimental, onMetrics) {
        this.llmClient = llmClient;
        this.defaultModelName = defaultModelName;
        this.defaultClientOptions = defaultClientOptions;
        this.resolveLlmClient = resolveLlmClient;
        this.systemPrompt = systemPrompt ?? "";
        this.logInferenceToFile = logInferenceToFile ?? false;
        this.experimental = experimental ?? false;
        this.onMetrics = onMetrics;
    }
    async extract(params) {
        const { instruction, schema, page, selector, timeout, model } = params;
        const llmClient = this.resolveLlmClient(model);
        const effectiveTimeoutMs = typeof timeout === "number" && timeout > 0 ? timeout : undefined;
        const ensureTimeRemaining = createTimeoutGuard(effectiveTimeoutMs, (ms) => new ExtractTimeoutError(ms));
        // No-args → page text (parity with v2)
        const noArgs = !instruction && !schema;
        if (noArgs) {
            const focusSelector = selector?.replace(/^xpath=/i, "") ?? "";
            ensureTimeRemaining();
            const snap = await captureHybridSnapshot(page, {
                experimental: this.experimental,
                focusSelector: focusSelector || undefined,
            });
            const result = { pageText: snap.combinedTree };
            // Validate via the same schema used in v2
            return pageTextSchema.parse(result);
        }
        if (!instruction && schema) {
            throw new StagehandInvalidArgumentError("extract() requires an instruction when a schema is provided.");
        }
        const focusSelector = selector?.replace(/^xpath=/, "") ?? "";
        // Build the hybrid snapshot (includes combinedTree; combinedUrlMap optional)
        ensureTimeRemaining();
        const { combinedTree, combinedUrlMap } = await captureHybridSnapshot(page, {
            experimental: this.experimental,
            focusSelector: focusSelector,
        });
        v3Logger({
            category: "extraction",
            message: "Starting extraction using a11y snapshot",
            level: 1,
            auxiliary: instruction
                ? { instruction: { value: instruction, type: "string" } }
                : undefined,
        });
        // Normalize schema: if instruction provided without schema, use defaultExtractSchema
        const baseSchema = (schema ??
            defaultExtractSchema);
        // Ensure we pass an object schema into inference; wrap non-object schemas
        const isObjectSchema = getZodType(baseSchema) === "object";
        const WRAP_KEY = "value";
        const factory = getZFactory(baseSchema);
        const objectSchema = isObjectSchema
            ? baseSchema
            : factory.object({
                [WRAP_KEY]: baseSchema,
            });
        const [transformedSchema, urlFieldPaths] = transformUrlStringsToNumericIds(objectSchema);
        ensureTimeRemaining();
        const extractionResponse = await runExtract({
            instruction,
            domElements: combinedTree,
            schema: transformedSchema,
            llmClient,
            userProvidedInstructions: this.systemPrompt,
            logger: v3Logger,
            logInferenceToFile: this.logInferenceToFile,
        });
        const { metadata: { completed }, prompt_tokens, completion_tokens, reasoning_tokens = 0, cached_input_tokens = 0, inference_time_ms, ...rest } = extractionResponse;
        let output = rest;
        // Update EXTRACT metrics from the LLM calls
        this.onMetrics?.(V3FunctionName.EXTRACT, prompt_tokens, completion_tokens, reasoning_tokens, cached_input_tokens, inference_time_ms);
        // Re-inject URLs for any url() fields we temporarily converted to number()
        const idToUrl = (combinedUrlMap ?? {});
        for (const { segments } of urlFieldPaths) {
            injectUrls(output, segments, idToUrl);
        }
        // If we wrapped a non-object schema, unwrap the value
        if (!isObjectSchema && output && typeof output === "object") {
            output = output[WRAP_KEY];
        }
        const resultPreviewLength = 200;
        const resultString = JSON.stringify(output) ?? "undefined";
        const resultPreview = resultString.length > resultPreviewLength
            ? resultString.slice(0, resultPreviewLength) + "..."
            : resultString;
        v3Logger({
            category: "extraction",
            message: completed
                ? "Extraction completed successfully"
                : "Extraction incomplete after processing all data",
            level: 1,
            auxiliary: {
                prompt_tokens: { value: String(prompt_tokens), type: "string" },
                completion_tokens: { value: String(completion_tokens), type: "string" },
                inference_time_ms: {
                    value: String(inference_time_ms),
                    type: "string",
                },
                result: { value: resultPreview, type: "string" },
            },
        });
        return output;
    }
}
//# sourceMappingURL=extractHandler.js.map