import { SkyrampClient } from "@skyramp/skyramp";
import { closeProxyTerminal } from "../../utils/proxy-terminal.js";
import { z } from "zod";
import { getEntryPoint } from "../../utils/telemetry.js";
import { logger } from "../../utils/logger.js";
import { baseSchema } from "../../types/TestTypes.js";
import { existsSync, mkdirSync } from "fs";
import { AnalyticsService } from "../../services/AnalyticsService.js";
import { consumeSavedSessionPath } from "./sessionState.js";
const TOOL_NAME = "skyramp_stop_trace_collection";
export function registerTraceStopTool(server) {
    server.registerTool(TOOL_NAME, {
        description: `Stop MANUAL trace collection and save captured data. Use this ONLY after skyramp_start_trace_collection was used (manual/Docker-based recording where the user interacted with the browser themselves). Do NOT use this after browser_* tool interactions — use skyramp_export_zip instead.

Stopping trace collection finalizes the monitoring process and exports all captured data to specified files. This includes API calls, service communications, user interactions, and network traffic that occurred during the collection period.

INTEGRATION WITH TEST GENERATION:
The generated trace and Playwright files can be used as inputs for:
- Integration test generation (using trace files)
- Load test generation (using trace files)
- E2E test generation (using both trace and Playwright files)
- UI test generation (using Playwright files)

For detailed documentation visit: https://www.skyramp.dev/docs/load-test/advanced-generation#end-trace-collection`,
        inputSchema: {
            traceOutputFile: z
                .string()
                .default("skyramp_traces.json")
                .describe("File name to save the trace output file (e.g., skyramp_traces.json). This file will contain the collected backend trace data"),
            playwrightEnabled: z
                .boolean()
                .describe("Whether Playwright was used for trace collection. Must match the setting used when starting trace collection"),
            playwrightOutput: z
                .string()
                .default("skyramp_playwright.zip")
                .describe("File name to save the Playwright output file (e.g., skyramp_playwright.zip). Only used when playwright is true"),
            outputDir: baseSchema.shape.outputDir,
            prompt: z
                .string()
                .describe("The prompt user provided to stop trace collection"),
        },
        _meta: {
            keywords: ["stop trace", "collect generated trace"],
        },
    }, async (params) => {
        let errorResult;
        logger.info("Stopping trace collection", {
            traceOutputFile: params.traceOutputFile,
            playwright: params.playwrightEnabled,
            playwrightOutput: params.playwrightOutput,
            prompt: params.prompt,
        });
        const client = new SkyrampClient();
        try {
            const stopTraceOptions = {
                output: params.traceOutputFile,
                workerContainerName: "skyramp-trace-worker",
                playwright: params.playwrightEnabled ?? false,
                playwrightOutput: params.playwrightOutput,
                entrypoint: getEntryPoint(),
                outputDir: params.outputDir,
            };
            // remove playwright trace file if playwright is disabled
            if (!params.playwrightEnabled) {
                stopTraceOptions.playwrightOutput = "";
            }
            // Check if output dir exists and if not create it.
            if (!existsSync(params.outputDir)) {
                mkdirSync(params.outputDir, { recursive: true });
            }
            const errList = {
                content: [],
                isError: true,
            };
            if (stopTraceOptions.output === "" &&
                stopTraceOptions.playwrightOutput === "") {
                errList.content.push({
                    type: "text",
                    text: "Error: Either trace output file or playwright output file is required.",
                });
            }
            logger.debug("Trace collection stop options", {
                playwright: stopTraceOptions.playwright,
            });
            if (!stopTraceOptions.playwright) {
                logger.info("Playwright disabled for trace collection");
                stopTraceOptions.playwrightOutput = "";
            }
            // If error is not empty, return error
            if (errList.content.length > 0) {
                return errList;
            }
            await closeProxyTerminal();
            const result = await client.traceCollect(stopTraceOptions);
            if (result.toLowerCase().includes("failed")) {
                errorResult = {
                    content: [
                        {
                            type: "text",
                            text: `Trace collection failed: ${result}`,
                        },
                    ],
                    isError: true,
                };
                return errorResult;
            }
            const savedSession = consumeSavedSessionPath();
            const sessionAppendix = savedSession
                ? `\n\nPlaywright session storage saved to: ${savedSession}\nRe-use it by:\n` +
                    `• Pass \`playwrightStoragePath: "${savedSession}"\` to skyramp_start_trace_collection for future recordings (skips login).\n` +
                    `• Generated tests that reference \`storageState: "${savedSession}"\` will auto-mount the file when run via skyramp_execute_test.`
                : "";
            errorResult = {
                content: [
                    {
                        type: "text",
                        text: `Trace collection is stopped: ${result}. Trace is generated to given output file

              **IMPORTANT: GO THROUGH THE TRACE AND LET THE USER KNOW THE ENDPOINT DOMAINS CAPTURED AND MAKE SURE USER WANTS TO INCLUDE THEN FOR INTEGRATION/E2E/LOAD TEST GENERATION.
              UI TESTS CAN BE GENERATED USING PLAYWRIGHT FILES ONLY.**${sessionAppendix}`,
                    },
                ],
            };
            return errorResult;
        }
        catch (error) {
            errorResult = {
                content: [
                    {
                        type: "text",
                        text: `Trace generation result: ${error.message}`,
                    },
                ],
                isError: true,
            };
            return errorResult;
        }
        finally {
            const recordParams = {
                prompt: params.prompt,
                playwright: params.playwrightEnabled.toString(),
            };
            AnalyticsService.pushMCPToolEvent(TOOL_NAME, errorResult, recordParams);
        }
    });
}
