/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { ClassifierStrategy } from './classifierStrategy.js';
import { isFunctionCall, isFunctionResponse, } from '../../utils/messageInspectors.js';
import { DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_MODEL_AUTO, PREVIEW_GEMINI_MODEL_AUTO, PREVIEW_GEMINI_3_1_MODEL, PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, } from '../../config/models.js';
import { promptIdContext } from '../../utils/promptIdContext.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { AuthType } from '../../core/contentGenerator.js';
import { ModelAvailabilityService } from '../../availability/modelAvailabilityService.js';
vi.mock('../../core/baseLlmClient.js');
describe('ClassifierStrategy', () => {
    let strategy;
    let mockContext;
    let mockConfig;
    let mockBaseLlmClient;
    let mockLocalLiteRtLmClient;
    let mockResolvedConfig;
    beforeEach(() => {
        vi.clearAllMocks();
        strategy = new ClassifierStrategy();
        mockContext = {
            history: [],
            request: [{ text: 'simple task' }],
            signal: new AbortController().signal,
        };
        mockResolvedConfig = {
            model: 'classifier',
            generateContentConfig: {},
        };
        mockConfig = {
            modelConfigService: {
                getResolvedConfig: vi.fn().mockReturnValue(mockResolvedConfig),
            },
            getModel: vi.fn().mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO),
            getNumericalRoutingEnabled: vi.fn().mockResolvedValue(false),
            getGemini31Launched: vi.fn().mockResolvedValue(false),
            getUseCustomToolModel: vi.fn().mockImplementation(async () => {
                const launched = await mockConfig.getGemini31Launched();
                const authType = mockConfig.getContentGeneratorConfig().authType;
                return launched && authType === AuthType.USE_GEMINI;
            }),
            getContentGeneratorConfig: vi.fn().mockReturnValue({
                authType: AuthType.LOGIN_WITH_GOOGLE,
            }),
            getModelAvailabilityService: vi
                .fn()
                .mockReturnValue(new ModelAvailabilityService()),
        };
        mockBaseLlmClient = {
            generateJson: vi.fn(),
        };
        mockLocalLiteRtLmClient = {};
        vi.spyOn(promptIdContext, 'getStore').mockReturnValue('test-prompt-id');
    });
    it('should return null if numerical routing is enabled and model is Gemini 3', async () => {
        vi.mocked(mockConfig.getNumericalRoutingEnabled).mockResolvedValue(true);
        vi.mocked(mockConfig.getModel).mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO);
        const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(decision).toBeNull();
        expect(mockBaseLlmClient.generateJson).not.toHaveBeenCalled();
    });
    it('should NOT return null if numerical routing is enabled but model is NOT Gemini 3', async () => {
        vi.mocked(mockConfig.getNumericalRoutingEnabled).mockResolvedValue(true);
        vi.mocked(mockConfig.getModel).mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO);
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({
            reasoning: 'test',
            model_choice: 'flash',
        });
        const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(decision).not.toBeNull();
        expect(mockBaseLlmClient.generateJson).toHaveBeenCalled();
    });
    it('should call generateJson with the correct parameters', async () => {
        const mockApiResponse = {
            reasoning: 'Simple task',
            model_choice: 'flash',
        };
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
        await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(mockBaseLlmClient.generateJson).toHaveBeenCalledWith(expect.objectContaining({
            modelConfigKey: { model: mockResolvedConfig.model },
            promptId: 'test-prompt-id',
        }));
    });
    it('should route to FLASH model for a simple task', async () => {
        const mockApiResponse = {
            reasoning: 'This is a simple task.',
            model_choice: 'flash',
        };
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
        const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(mockBaseLlmClient.generateJson).toHaveBeenCalledOnce();
        expect(decision).toEqual({
            model: DEFAULT_GEMINI_FLASH_MODEL,
            metadata: {
                source: 'Classifier',
                latencyMs: expect.any(Number),
                reasoning: mockApiResponse.reasoning,
            },
        });
    });
    it('should route to PRO model for a complex task', async () => {
        const mockApiResponse = {
            reasoning: 'This is a complex task.',
            model_choice: 'pro',
        };
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
        mockContext.request = [{ text: 'how do I build a spaceship?' }];
        const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(mockBaseLlmClient.generateJson).toHaveBeenCalledOnce();
        expect(decision).toEqual({
            model: DEFAULT_GEMINI_MODEL,
            metadata: {
                source: 'Classifier',
                latencyMs: expect.any(Number),
                reasoning: mockApiResponse.reasoning,
            },
        });
    });
    it('should return null if the classifier API call fails', async () => {
        const consoleWarnSpy = vi
            .spyOn(debugLogger, 'warn')
            .mockImplementation(() => { });
        const testError = new Error('API Failure');
        vi.mocked(mockBaseLlmClient.generateJson).mockRejectedValue(testError);
        const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(decision).toBeNull();
        expect(consoleWarnSpy).toHaveBeenCalled();
        consoleWarnSpy.mockRestore();
    });
    it('should return null if the classifier returns a malformed JSON object', async () => {
        const consoleWarnSpy = vi
            .spyOn(debugLogger, 'warn')
            .mockImplementation(() => { });
        const malformedApiResponse = {
            reasoning: 'This is a simple task.',
            // model_choice is missing, which will cause a Zod parsing error.
        };
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(malformedApiResponse);
        const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(decision).toBeNull();
        expect(consoleWarnSpy).toHaveBeenCalled();
        consoleWarnSpy.mockRestore();
    });
    it('should filter out tool-related history before sending to classifier', async () => {
        mockContext.history = [
            { role: 'user', parts: [{ text: 'call a tool' }] },
            { role: 'model', parts: [{ functionCall: { name: 'test_tool' } }] },
            {
                role: 'user',
                parts: [
                    { functionResponse: { name: 'test_tool', response: { ok: true } } },
                ],
            },
            { role: 'user', parts: [{ text: 'another user turn' }] },
        ];
        const mockApiResponse = {
            reasoning: 'Simple.',
            model_choice: 'flash',
        };
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
        await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        const generateJsonCall = vi.mocked(mockBaseLlmClient.generateJson).mock
            .calls[0][0];
        const contents = generateJsonCall.contents;
        const expectedContents = [
            { role: 'user', parts: [{ text: 'call a tool' }] },
            { role: 'user', parts: [{ text: 'another user turn' }] },
            { role: 'user', parts: [{ text: 'simple task' }] },
        ];
        expect(contents).toEqual(expectedContents);
    });
    it('should respect HISTORY_SEARCH_WINDOW and HISTORY_TURNS_FOR_CONTEXT', async () => {
        const longHistory = [];
        for (let i = 0; i < 30; i++) {
            longHistory.push({ role: 'user', parts: [{ text: `Message ${i}` }] });
            // Add noise that should be filtered
            if (i % 2 === 0) {
                longHistory.push({
                    role: 'model',
                    parts: [{ functionCall: { name: 'noise', args: {} } }],
                });
            }
        }
        mockContext.history = longHistory;
        const mockApiResponse = {
            reasoning: 'Simple.',
            model_choice: 'flash',
        };
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
        await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        const generateJsonCall = vi.mocked(mockBaseLlmClient.generateJson).mock
            .calls[0][0];
        const contents = generateJsonCall.contents;
        // Manually calculate what the history should be
        const HISTORY_SEARCH_WINDOW = 20;
        const HISTORY_TURNS_FOR_CONTEXT = 4;
        const historySlice = longHistory.slice(-HISTORY_SEARCH_WINDOW);
        const cleanHistory = historySlice.filter((content) => !isFunctionCall(content) && !isFunctionResponse(content));
        const finalHistory = cleanHistory.slice(-HISTORY_TURNS_FOR_CONTEXT);
        expect(contents).toEqual([
            ...finalHistory,
            { role: 'user', parts: mockContext.request },
        ]);
        // There should be 4 history items + the current request
        expect(contents).toHaveLength(5);
    });
    it('should use a fallback promptId if not found in context', async () => {
        const consoleWarnSpy = vi
            .spyOn(debugLogger, 'warn')
            .mockImplementation(() => { });
        vi.spyOn(promptIdContext, 'getStore').mockReturnValue(undefined);
        const mockApiResponse = {
            reasoning: 'Simple.',
            model_choice: 'flash',
        };
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
        await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        const generateJsonCall = vi.mocked(mockBaseLlmClient.generateJson).mock
            .calls[0][0];
        expect(generateJsonCall.promptId).toMatch(/^classifier-router-fallback-\d+-\w+$/);
        expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining('Could not find promptId in context for classifier-router. This is unexpected. Using a fallback ID:'));
        consoleWarnSpy.mockRestore();
    });
    it('should respect requestedModel from context in resolveClassifierModel', async () => {
        const requestedModel = DEFAULT_GEMINI_MODEL; // Pro model
        const mockApiResponse = {
            reasoning: 'Choice is flash',
            model_choice: 'flash',
        };
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
        const contextWithRequestedModel = {
            ...mockContext,
            requestedModel,
        };
        const decision = await strategy.route(contextWithRequestedModel, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(decision).not.toBeNull();
        // Since requestedModel is Pro, and choice is flash, it should resolve to Flash
        expect(decision?.model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
    });
    it('should return null (bypass classifier) if history is only tool turns and request is a function response', async () => {
        const history = [
            { role: 'model', parts: [{ functionCall: { name: 'tool' } }] },
            {
                role: 'user',
                parts: [{ functionResponse: { name: 'tool', response: { ok: true } } }],
            },
            { role: 'model', parts: [{ functionCall: { name: 'tool2' } }] },
        ];
        mockContext.history = history;
        mockContext.request = [
            { functionResponse: { name: 'tool2', response: { ok: true } } },
        ];
        const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(decision).toBeNull();
        expect(mockBaseLlmClient.generateJson).not.toHaveBeenCalled();
    });
    it('should return null (bypass classifier) if history has text turns and request is a function response', async () => {
        const history = [
            { role: 'user', parts: [{ text: 'some task' }] },
            { role: 'model', parts: [{ functionCall: { name: 'tool' } }] },
        ];
        mockContext.history = history;
        mockContext.request = [
            { functionResponse: { name: 'tool', response: { ok: true } } },
        ];
        const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(decision).toBeNull();
        expect(mockBaseLlmClient.generateJson).not.toHaveBeenCalled();
    });
    it('should still route if history is only tool turns but request is text', async () => {
        const history = [
            { role: 'model', parts: [{ functionCall: { name: 'tool' } }] },
            {
                role: 'user',
                parts: [{ functionResponse: { name: 'tool', response: { ok: true } } }],
            },
            { role: 'model', parts: [{ functionCall: { name: 'tool2' } }] },
        ];
        mockContext.history = history;
        mockContext.request = [{ text: 'simple task' }];
        const mockApiResponse = {
            reasoning: 'Simple.',
            model_choice: 'flash',
        };
        vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
        const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
        expect(decision).not.toBeNull();
        expect(mockBaseLlmClient.generateJson).toHaveBeenCalled();
        const generateJsonCall = vi.mocked(mockBaseLlmClient.generateJson).mock
            .calls[0][0];
        const contents = generateJsonCall.contents;
        // History should be empty because all turns were tool turns and stripped.
        // Request should be present.
        const expectedContents = [
            {
                role: 'user',
                parts: [{ text: 'simple task' }],
            },
        ];
        expect(contents).toEqual(expectedContents);
    });
    describe('Gemini 3.1 and Custom Tools Routing', () => {
        it('should route to PREVIEW_GEMINI_3_1_MODEL when Gemini 3.1 is launched', async () => {
            vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
            vi.mocked(mockConfig.getModel).mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO);
            const mockApiResponse = {
                reasoning: 'Complex task',
                model_choice: 'pro',
            };
            vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
            const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
            expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_MODEL);
        });
        it('should route to PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL when Gemini 3.1 is launched and auth is USE_GEMINI', async () => {
            vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
            vi.mocked(mockConfig.getModel).mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO);
            vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
                authType: AuthType.USE_GEMINI,
            });
            const mockApiResponse = {
                reasoning: 'Complex task',
                model_choice: 'pro',
            };
            vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
            const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
            expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
        });
        it('should route to DEFAULT_GEMINI_FLASH_MODEL when hasGemini35FlashGAAccess is true', async () => {
            mockConfig.hasGemini35FlashGAAccess = vi.fn().mockReturnValue(true);
            vi.mocked(mockConfig.getModel).mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO);
            const mockApiResponse = {
                reasoning: 'Simple task',
                model_choice: 'flash',
            };
            vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(mockApiResponse);
            const decision = await strategy.route(mockContext, mockConfig, mockBaseLlmClient, mockLocalLiteRtLmClient);
            expect(decision?.model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
        });
    });
});
//# sourceMappingURL=classifierStrategy.test.js.map