import { z } from "zod";
import fs from "fs";
import { logger } from "../../utils/logger.js";
import { ProgrammingLanguage, TestType } from "../../types/TestTypes.js";
import { ModularizationService, } from "../../services/ModularizationService.js";
import { AnalyticsService } from "../../services/AnalyticsService.js";
import { normalizeLanguageParams, resolveParamAliases, } from "../../utils/normalizeParams.js";
import { normalizeSkyrampImportsInFile } from "../../utils/normalizeSkyrampImports.js";
import { stageGeneratedPaths } from "../../utils/gitStaging.js";
const modularizationSchema = {
    testFile: z
        .string()
        .describe("The test file to process with modularization principles applied"),
    language: z
        .nativeEnum(ProgrammingLanguage)
        .optional()
        .describe("The programming language of the test file. Inferred from file extension if not provided."),
    testType: z
        .enum([TestType.UI, TestType.E2E, TestType.INTEGRATION])
        .optional()
        .describe("Specifies the type of test (UI, E2E, or Integration). Defaults to 'ui'. DO NOT USE TEST NAME/DESCRIPTION AS TEST TYPE."),
    isTraceBased: z
        .boolean()
        .default(false)
        .describe("Whether the test is generated from traces. MUST be set to true when the test was generated from a Playwright trace (browser_* tools + skyramp_export_zip) or from skyramp_start/stop_trace_collection."),
    prompt: z
        .string()
        .optional()
        .describe("The prompt or code content to process. If not provided, the test file content is used."),
};
const TOOL_NAME = "skyramp_modularization";
export function registerModularizationTool(server) {
    server.registerTool(TOOL_NAME, {
        description: `Provides modularization instructions and test file content. YOU MUST immediately modularize the code and write it back to the file. The tool returns:
1. Complete modularization instructions
2. The current test file content
3. Step-by-step requirements you must complete

**MANDATORY**: ONLY MODULARIZE THE TEST IF THE IS TRACE BASED FLAG IS SET TO TRUE ELSE DO NOT MODULARIZE THE TEST AND LEAVE THE TEST FILE AS IS.
**CRITICAL**: NON TRACE BASED TESTS ARE ALREADY MODULARIZED AND DO NOT NEED MODULARIZATION.
After modularization, if errors remain, call skyramp_fix_errors.

**CRITICAL**: This tool expects you to complete the modularization by writing the refactored code back to the file.`,
        inputSchema: modularizationSchema,
        _meta: {
            keywords: [
                "modularization",
                "code organization",
                "structure",
                "refactor",
            ],
        },
    }, async (params) => {
        let errorResult;
        try {
            // Resolve aliases and normalize
            resolveParamAliases(params, {
                testFilePath: "testFile",
                file: "testFile",
            });
            normalizeLanguageParams(params);
            // Infer language from file extension if not provided
            if (!params.language && params.testFile) {
                if (params.testFile.endsWith(".ts"))
                    params.language = "typescript";
                else if (params.testFile.endsWith(".js"))
                    params.language = "javascript";
                else if (params.testFile.endsWith(".py"))
                    params.language = "python";
                else if (params.testFile.endsWith(".java"))
                    params.language = "java";
                else
                    params.language = "typescript";
            }
            // Default testType to 'ui'
            if (!params.testType)
                params.testType = TestType.UI;
            // For UI/E2E/Integration tests, always enable modularization.
            // These test types are trace-based and benefit from modularization.
            if (!params.isTraceBased && [TestType.UI, TestType.E2E, TestType.INTEGRATION].includes(params.testType))
                params.isTraceBased = true;
            normalizeSkyrampImportsInFile(params.testFile);
            // Stage so testbot includes the generated files in its output commit.
            await stageGeneratedPaths(params.testFile);
            // Default prompt to test file content
            if (!params.prompt && params.testFile) {
                try {
                    params.prompt = fs.readFileSync(params.testFile, "utf8");
                }
                catch {
                    params.prompt = `Modularize the test in ${params.testFile}`;
                }
            }
            logger.info("Generating modularization", {
                testFile: params.testFile,
                language: params.language,
                prompt: params.prompt?.substring(0, 100),
                testType: params.testType,
            });
            const service = new ModularizationService();
            return await service.processModularizationRequest(params);
        }
        catch (error) {
            const errorMessage = `Modularization failed: ${error.message}`;
            logger.error(errorMessage, error);
            errorResult = {
                content: [
                    {
                        type: "text",
                        text: errorMessage,
                    },
                ],
                isError: true,
            };
            return errorResult;
        }
        finally {
            AnalyticsService.pushMCPToolEvent(TOOL_NAME, errorResult, {
                prompt: params.prompt ? `[${params.prompt.length} chars]` : "",
                testFile: params.testFile,
                language: params.language ?? "",
                testType: params.testType ?? "",
            });
        }
    });
}
