import { z } from "zod";
import fs from "fs";
import path from "path";
import { SkyrampClient } from "@skyramp/skyramp";
import openProxyTerminalTracked from "../../utils/proxy-terminal.js";
import { getEntryPoint } from "../../utils/telemetry.js";
import { logger } from "../../utils/logger.js";
import { basePlaywrightSchema, baseSchema, SESSION_STORAGE_FILENAME, } from "../../types/TestTypes.js";
import { AnalyticsService } from "../../services/AnalyticsService.js";
import { resolveSessionPaths } from "./resolveSessionPaths.js";
import { setSavedSessionPath } from "./sessionState.js";
const TOOL_NAME = "skyramp_start_trace_collection";
export function registerTraceTool(server) {
    server.registerTool(TOOL_NAME, {
        description: `Start MANUAL trace collection using Skyramp's Docker-based tracing.

Use this tool ONLY when the user explicitly asks for "manual recording", "manual trace", "record myself", "I will interact", "Docker trace", or wants to interact with the browser themselves. For AI-driven browser automation (where the agent fills forms, clicks buttons, etc.), use the browser_* tools instead.

This tool launches a Docker proxy that captures network traffic while the USER manually interacts with the application in a browser window. The user performs clicks, form fills, and navigation themselves.

WORKFLOW:
1. Start trace collection — opens a browser for the user
2. USER manually interacts with the application
3. Stop trace collection to save captured data
4. Use traces to generate test scenarios

SESSION HANDLING:
Pass \`sessionMode\` to declare what you want to do with the workspace's session file. Defaults to \`auto\`, which does the right thing for most prompts.

\`sessionMode: "capture"\` — Use when the user wants to log in once and SAVE a session for reuse later. Triggers: "save my session", "log in once", "store the session", "create a new session", "capture login".
  Workflow:
  1. Start trace collection (sessionMode=capture).
  2. USER logs in.
  3. Stop trace collection AFTER the post-login page is fully loaded — wait a few seconds for the app's initial authenticated API calls to fire, then stop. Do not interact further; those actions would pollute the trace and dilute the saved cookies.
  4. The session file lands at \`outputDir/${SESSION_STORAGE_FILENAME}\`.

\`sessionMode: "reuse"\` — Use when the user wants to RECORD a test flow against an existing authenticated session. Triggers: "load my session", "use my saved session", "skip login", "reuse the session", "I'm already logged in".
  Workflow:
  1. Start trace collection (sessionMode=reuse). The recorder loads \`outputDir/${SESSION_STORAGE_FILENAME}\` so the browser starts already authenticated.
  2. USER walks through the test flow (no login needed).
  3. Stop trace collection when the flow is complete.
  4. The session file is NOT overwritten — this preserves the captured session for the next recording or test run.

\`sessionMode: "ignore"\` — Use when the user explicitly does NOT want session handling. Triggers: "fresh trace", "without session", "ignore my session".

\`sessionMode: "auto"\` (default) — Decides between capture and reuse based on whether a session file already exists in \`outputDir\`. Use when the user does not signal an intent either way. If a session file is present, behaves like \`reuse\`; otherwise like \`capture\`.

For detailed documentation visit: https://www.skyramp.dev/docs/load-test/advanced-generation#start-trace-collection`,
        inputSchema: {
            playwright: z
                .boolean()
                .describe("Whether to enable Playwright for trace collection. Set to true for UI interactions, false for API-only tracing")
                .default(true),
            sessionMode: z
                .enum(["auto", "capture", "reuse", "ignore"])
                .default("auto")
                .describe(`Controls how Playwright session storage is handled for this trace. ` +
                `"capture" = save a fresh session at outputDir/${SESSION_STORAGE_FILENAME} (use when user wants to log in once and save). ` +
                `"reuse" = load that file into the browser and DO NOT overwrite it (use when user wants to record a test flow against an existing session). ` +
                `"ignore" = neither load nor save. ` +
                `"auto" (default) = reuse if the session file already exists in outputDir, otherwise capture. ` +
                `Explicit playwrightStoragePath / playwrightSaveStoragePath overrides still win when provided.`),
            browser: basePlaywrightSchema.shape.browser,
            device: basePlaywrightSchema.shape.device,
            playwrightStoragePath: basePlaywrightSchema.shape.playwrightStoragePath,
            playwrightSaveStoragePath: basePlaywrightSchema.shape.playwrightSaveStoragePath,
            playwrightViewportSize: basePlaywrightSchema.shape.playwrightViewportSize,
            runtime: z
                .string()
                .default("docker")
                .describe("Runtime environment for trace collection. Currently only 'docker' is supported and is used as the default."),
            include: z
                .array(z.string())
                .default([])
                .describe("List of service names or patterns to include in trace collection"),
            exclude: z
                .array(z.string())
                .default([])
                .describe("List of service names or patterns to exclude from trace collection"),
            noProxy: z
                .array(z.string())
                .default([])
                .describe("List of hosts or patterns that should bypass proxy during tracing"),
            dockerNetwork: z
                .string()
                .default("skyramp")
                .describe("Docker network name to use for containerized trace collection"),
            dockerWorkerPort: z
                .number()
                .default(35142)
                .describe("Port number for the Docker worker service during trace collection"),
            outputDir: baseSchema.shape.outputDir,
            prompt: z
                .string()
                .describe("The prompt user provided to start trace collection"),
        },
        _meta: {
            keywords: ["start trace", "trace collection", "trace generation"],
        },
    }, async (params, extra) => {
        let errorResult;
        // Helper to send progress notifications to the MCP client
        const sendProgress = async (progress, total, message) => {
            const progressToken = extra._meta?.progressToken;
            if (progressToken !== undefined) {
                const notification = {
                    method: "notifications/progress",
                    params: {
                        progressToken,
                        progress,
                        total,
                        message,
                    },
                };
                await extra.sendNotification(notification);
            }
        };
        logger.info("Starting trace collection", {
            playwright: params.playwright,
            browser: params.browser,
            device: params.device,
            runtime: params.runtime,
            include: params.include,
            exclude: params.exclude,
            noProxy: params.noProxy,
            dockerNetwork: params.dockerNetwork,
            dockerWorkerPort: params.dockerWorkerPort,
            outputDir: params.outputDir,
            prompt: params.prompt,
        });
        const defaultSessionFile = path.join(params.outputDir, SESSION_STORAGE_FILENAME);
        const sessionExists = (() => {
            try {
                return fs.existsSync(defaultSessionFile);
            }
            catch {
                return false;
            }
        })();
        const { loadPath, savePath } = resolveSessionPaths({
            mode: params.sessionMode,
            loadOverride: params.playwrightStoragePath,
            saveOverride: params.playwrightSaveStoragePath,
            outputDir: params.outputDir,
            sessionExists,
        });
        logger.info("Resolved session paths", {
            sessionMode: params.sessionMode,
            sessionExists,
            loadPath: loadPath ?? "(not loading)",
            savePath: savePath ?? "(not saving)",
        });
        // Carry the save path forward so the stop tool can name it in its message.
        // When savePath is undefined (e.g. reuse mode), the stop tool just won't
        // reference one — which is correct: the session file isn't being touched.
        setSavedSessionPath(savePath);
        try {
            // Send initial progress
            await sendProgress(0, 100, "Initializing trace collection...");
            const client = new SkyrampClient();
            const generateOptions = {
                testType: "trace",
                runtime: params.runtime,
                dockerNetwork: params.dockerNetwork,
                dockerWorkerPort: params.dockerWorkerPort.toString(),
                generateInclude: params.include,
                generateExclude: params.exclude,
                generateNoProxy: params.noProxy,
                unblock: true,
                playwright: params.playwright,
                browser: params.browser,
                device: params.device,
                playwrightViewportSize: params.playwrightViewportSize,
                entrypoint: getEntryPoint(),
            };
            // Only pass the storage options when sessionMode actually wants them.
            // Leaving them undefined tells the underlying recorder to skip load /
            // save entirely — that's what prevents the silent-overwrite trap.
            if (loadPath)
                generateOptions.playwrightStoragePath = loadPath;
            if (savePath)
                generateOptions.playwrightSaveStoragePath = savePath;
            // Send progress for configuration
            const traceMode = params.playwright ? "UI + Backend" : "Backend-only";
            await sendProgress(30, 100, `Configuring ${traceMode} trace collection...`);
            // Start periodic progress updates during the long-running operation
            const startTime = Date.now();
            let progressTick = 0;
            const progressInterval = setInterval(() => {
                progressTick++;
                const elapsed = Math.floor((Date.now() - startTime) / 1000);
                // Progress moves from 30% to 65% during service startup
                const percent = Math.min(30 + progressTick * 5, 65);
                sendProgress(percent, 100, `Starting trace collection service... (${elapsed}s elapsed)`).catch(() => {
                    // Ignore progress notification errors
                });
                // Also log to server for visibility
                logger.info(`Trace collection progress: ${percent}% (${elapsed}s elapsed)`);
            }, 10000); // 10 seconds
            let result;
            try {
                result = await client.generateRestTest(generateOptions);
            }
            finally {
                clearInterval(progressInterval);
            }
            if (result.toLowerCase().includes("failed")) {
                // Clear stashed session path so a failed start does not leak into a later unrelated stop.
                setSavedSessionPath(undefined);
                errorResult = {
                    content: [
                        {
                            type: "text",
                            text: `Trace collection failed: ${result}`,
                        },
                    ],
                    isError: true,
                };
                return errorResult;
            }
            // Send progress for opening proxy terminal
            await sendProgress(70, 100, "Opening proxy terminal...");
            await openProxyTerminalTracked();
            // Send completion progress
            await sendProgress(100, 100, "Trace collection started successfully");
            const sessionGuidance = (() => {
                if (loadPath && savePath) {
                    return `\n\nSession ${loadPath} loaded into the browser; will be re-saved to ${savePath} on stop.`;
                }
                if (loadPath) {
                    return `\n\nSession ${loadPath} loaded into the browser — the browser starts already authenticated. The session file will NOT be overwritten on stop.`;
                }
                if (savePath) {
                    return `\n\nPlaywright session storage will be saved to ${savePath}. If the goal is to capture an authenticated session, ask the user to log in and then call skyramp_stop_trace_collection a few seconds AFTER the post-login page is fully loaded — long enough for the app's initial authenticated API calls to fire, short enough that no test-flow interactions are captured.`;
                }
                return "";
            })();
            errorResult = {
                content: [
                    {
                        type: "text",
                        text: `Trace collection started: ${result}. Please let me know when you are ready to stop the trace collection.${sessionGuidance}`,
                    },
                ],
            };
            return errorResult;
        }
        catch (error) {
            // Clear stashed session path so a failed start does not leak into a later unrelated stop.
            setSavedSessionPath(undefined);
            errorResult = {
                content: [
                    {
                        type: "text",
                        text: `Trace generation failed: ${error.message}`,
                    },
                ],
                isError: true,
            };
            return errorResult;
        }
        finally {
            const recordParams = {
                prompt: params.prompt,
                playwright: params.playwright.toString(),
            };
            AnalyticsService.pushMCPToolEvent(TOOL_NAME, errorResult, recordParams);
        }
    });
}
