/**
 * Unit tests for Bedrock Prompt Caching functionality
 *
 * Tests cover:
 * 1. CustomChatBedrockConverse - Tool caching with cachePoint via invocationParams
 * 2. AgentContext - System message format per provider:
 *    - Bedrock: plain string (cachePoint is added at the tool level by invocationParams)
 *    - Anthropic: cache_control added inline by AgentContext
 *    - OpenAI: plain string (no caching)
 */

import { CustomChatBedrockConverse } from '../index';
import { AgentContext } from '@/agents/AgentContext';
import { Providers } from '@/common';
import type * as t from '@/types';

type InvocationParamsOptions = Parameters<
  InstanceType<typeof CustomChatBedrockConverse>['invocationParams']
>[0];

describe('Bedrock Prompt Caching', () => {
  describe('CustomChatBedrockConverse - Tool Caching', () => {
    describe('invocationParams with promptCache enabled', () => {
      it('should add cachePoint to tools array when promptCache is true', () => {
        const model = new CustomChatBedrockConverse({
          model: 'anthropic.claude-3-haiku-20240307-v1:0',
          region: 'us-east-1',
          promptCache: true,
        });

        const mockTools = [
          {
            toolSpec: {
              name: 'get_weather',
              description: 'Get weather for a location',
              inputSchema: { json: { type: 'object', properties: {} } },
            },
          },
          {
            toolSpec: {
              name: 'search_web',
              description: 'Search the web',
              inputSchema: { json: { type: 'object', properties: {} } },
            },
          },
        ];

        const params = model.invocationParams({
          tools: mockTools,
        } satisfies InvocationParamsOptions);

        // Should have tools + cachePoint
        expect(params.toolConfig?.tools).toHaveLength(3);
        expect(params.toolConfig?.tools?.[2]).toEqual({
          cachePoint: { type: 'default' },
        });
      });

      it('should NOT add cachePoint when promptCache is false', () => {
        const model = new CustomChatBedrockConverse({
          model: 'anthropic.claude-3-haiku-20240307-v1:0',
          region: 'us-east-1',
          promptCache: false,
        });

        const mockTools = [
          {
            toolSpec: {
              name: 'get_weather',
              description: 'Get weather',
              inputSchema: { json: { type: 'object', properties: {} } },
            },
          },
        ];

        const params = model.invocationParams({
          tools: mockTools,
        } satisfies InvocationParamsOptions);

        // Should only have original tools, no cachePoint
        expect(params.toolConfig?.tools).toHaveLength(1);
        expect(params.toolConfig?.tools?.[0]).toEqual(mockTools[0]);
      });

      it('should NOT add cachePoint when promptCache is undefined (default)', () => {
        const model = new CustomChatBedrockConverse({
          model: 'anthropic.claude-3-haiku-20240307-v1:0',
          region: 'us-east-1',
        });

        const mockTools = [
          {
            toolSpec: {
              name: 'get_weather',
              description: 'Get weather',
              inputSchema: { json: { type: 'object', properties: {} } },
            },
          },
        ];

        const params = model.invocationParams({
          tools: mockTools,
        } satisfies InvocationParamsOptions);

        // Should only have original tools
        expect(params.toolConfig?.tools).toHaveLength(1);
      });

      it('should NOT add cachePoint when no tools are provided', () => {
        const model = new CustomChatBedrockConverse({
          model: 'anthropic.claude-3-haiku-20240307-v1:0',
          region: 'us-east-1',
          promptCache: true,
        });

        const params = model.invocationParams(
          {} satisfies InvocationParamsOptions
        );

        // toolConfig should be undefined or have no tools
        expect(params.toolConfig?.tools).toBeUndefined();
      });

      it('should NOT add cachePoint when tools array is empty', () => {
        const model = new CustomChatBedrockConverse({
          model: 'anthropic.claude-3-haiku-20240307-v1:0',
          region: 'us-east-1',
          promptCache: true,
        });

        const params = model.invocationParams({
          tools: [],
        } satisfies InvocationParamsOptions);

        // Empty tools array results in undefined toolConfig
        expect(params.toolConfig).toBeUndefined();
      });

      it('should preserve other invocationParams properties', () => {
        const model = new CustomChatBedrockConverse({
          model: 'anthropic.claude-3-haiku-20240307-v1:0',
          region: 'us-east-1',
          promptCache: true,
          temperature: 0.7,
          maxTokens: 1000,
        });

        const mockTools = [
          {
            toolSpec: {
              name: 'test_tool',
              description: 'Test',
              inputSchema: { json: { type: 'object', properties: {} } },
            },
          },
        ];

        const params = model.invocationParams({
          tools: mockTools,
        } satisfies InvocationParamsOptions);

        // Check that other params are preserved
        expect(params.inferenceConfig?.temperature).toBe(0.7);
        expect(params.inferenceConfig?.maxTokens).toBe(1000);
        // And cachePoint is still added
        expect(params.toolConfig?.tools).toHaveLength(2);
      });
    });

    describe('promptCache property', () => {
      it('should store promptCache value from constructor', () => {
        const modelWithCache = new CustomChatBedrockConverse({
          model: 'test-model',
          region: 'us-east-1',
          promptCache: true,
        });

        const modelWithoutCache = new CustomChatBedrockConverse({
          model: 'test-model',
          region: 'us-east-1',
          promptCache: false,
        });

        expect(modelWithCache.promptCache).toBe(true);
        expect(modelWithoutCache.promptCache).toBe(false);
      });

      it('should default promptCache to false when not provided', () => {
        const model = new CustomChatBedrockConverse({
          model: 'test-model',
          region: 'us-east-1',
        });

        expect(model.promptCache).toBe(false);
      });
    });
  });

  describe('AgentContext - System Message Caching', () => {
    const createBedrockContext = (options: {
      instructions?: string;
      promptCache?: boolean;
      additionalInstructions?: string;
    }): AgentContext => {
      const clientOptions: t.BedrockAnthropicInput = {
        model: 'anthropic.claude-3-haiku-20240307-v1:0',
        region: 'us-east-1',
        promptCache: options.promptCache,
      };

      return AgentContext.fromConfig({
        agentId: 'test-bedrock-agent',
        provider: Providers.BEDROCK,
        instructions: options.instructions,
        additional_instructions: options.additionalInstructions,
        clientOptions,
      });
    };

    const createAnthropicContext = (options: {
      instructions?: string;
      promptCachingEnabled?: boolean;
    }): AgentContext => {
      const clientOptions: t.AnthropicClientOptions = {
        promptCache: options.promptCachingEnabled,
      };
      return AgentContext.fromConfig({
        agentId: 'test-anthropic-agent',
        provider: Providers.ANTHROPIC,
        instructions: options.instructions,
        clientOptions,
      });
    };

    describe('Bedrock system message with promptCache: true', () => {
      it('should produce plain string system message (cachePoint is handled at tool level)', async () => {
        const ctx = createBedrockContext({
          instructions: 'You are a helpful assistant.',
          promptCache: true,
        });

        const systemRunnable = ctx.systemRunnable;
        expect(systemRunnable).toBeDefined();

        // Invoke the runnable to get the messages
        const result = await systemRunnable!.invoke([]);
        const systemMessage = result[0];

        // Bedrock system messages are plain strings — cachePoint is added
        // at the tool level by IllumaBedrockConverse.invocationParams()
        expect(typeof systemMessage.content).toBe('string');
        expect(systemMessage.content).toBe('You are a helpful assistant.');
      });

      it('should include combined instructions and additional_instructions as plain string', async () => {
        const ctx = createBedrockContext({
          instructions: 'Base instructions.',
          additionalInstructions: 'Additional context.',
          promptCache: true,
        });

        const result = await ctx.systemRunnable!.invoke([]);
        const systemMessage = result[0];

        // Bedrock system messages are plain strings even with promptCache
        expect(typeof systemMessage.content).toBe('string');
        expect(systemMessage.content).toContain('Base instructions.');
        expect(systemMessage.content).toContain('Additional context.');
      });
    });

    describe('Bedrock system message with promptCache: false', () => {
      it('should NOT add cachePoint when promptCache is false', async () => {
        const ctx = createBedrockContext({
          instructions: 'You are a helpful assistant.',
          promptCache: false,
        });

        const result = await ctx.systemRunnable!.invoke([]);
        const systemMessage = result[0];

        // Content should be plain string, not array with cachePoint
        expect(typeof systemMessage.content).toBe('string');
        expect(systemMessage.content).toBe('You are a helpful assistant.');
      });

      it('should NOT add cachePoint when promptCache is undefined', async () => {
        const ctx = createBedrockContext({
          instructions: 'You are a helpful assistant.',
          promptCache: undefined,
        });

        const result = await ctx.systemRunnable!.invoke([]);
        const systemMessage = result[0];

        expect(typeof systemMessage.content).toBe('string');
      });
    });

    describe('Anthropic system message caching (for comparison)', () => {
      it('should add cache_control for Anthropic with prompt-caching beta', async () => {
        const ctx = createAnthropicContext({
          instructions: 'You are a helpful assistant.',
          promptCachingEnabled: true,
        });

        const result = await ctx.systemRunnable!.invoke([]);
        const systemMessage = result[0];
        const content = systemMessage.content as Array<Record<string, unknown>>;

        expect(content).toHaveLength(1);
        expect(content[0]).toEqual({
          type: 'text',
          text: 'You are a helpful assistant.',
          cache_control: { type: 'ephemeral' },
        });
      });

      it('should NOT add cache_control for Anthropic without beta header', async () => {
        const ctx = createAnthropicContext({
          instructions: 'You are a helpful assistant.',
          promptCachingEnabled: false,
        });

        const result = await ctx.systemRunnable!.invoke([]);
        const systemMessage = result[0];

        expect(typeof systemMessage.content).toBe('string');
      });
    });

    describe('Provider-specific caching behavior', () => {
      it('should produce plain string for Bedrock (cachePoint handled at tool level)', async () => {
        const bedrockCtx = createBedrockContext({
          instructions: 'Test',
          promptCache: true,
        });

        const result = await bedrockCtx.systemRunnable!.invoke([]);

        // Bedrock system messages are plain strings — cachePoint is added
        // at the tool level by IllumaBedrockConverse.invocationParams(),
        // not inline in the system message content
        expect(typeof result[0].content).toBe('string');
        expect(result[0].content).toBe('Test');
      });

      it('should use cache_control format for Anthropic, not cachePoint', async () => {
        const anthropicCtx = createAnthropicContext({
          instructions: 'Test',
          promptCachingEnabled: true,
        });

        const result = await anthropicCtx.systemRunnable!.invoke([]);
        const content = result[0].content as Array<Record<string, unknown>>;

        // Anthropic uses cache_control inline in system message content
        expect(content.some((c) => 'cache_control' in c)).toBe(true);
        expect(content.some((c) => 'cachePoint' in c)).toBe(false);
      });

      it('should not add any caching for OpenAI provider', async () => {
        const openaiCtx = AgentContext.fromConfig({
          agentId: 'test-openai-agent',
          provider: Providers.OPENAI,
          instructions: 'Test instructions',
        });

        const result = await openaiCtx.systemRunnable!.invoke([]);
        const systemMessage = result[0];

        // OpenAI should have plain string content
        expect(typeof systemMessage.content).toBe('string');
        expect(systemMessage.content).toBe('Test instructions');
      });
    });
  });

  describe('Integration scenarios', () => {
    it('should handle multi-turn caching scenario', async () => {
      // Simulate: Query 1 with tools → Query 2 with same tools
      // Both should have cachePoint, enabling cache reuse

      const model = new CustomChatBedrockConverse({
        model: 'anthropic.claude-3-haiku-20240307-v1:0',
        region: 'us-east-1',
        promptCache: true,
      });

      const tools = [
        {
          toolSpec: {
            name: 'execute_code',
            description: 'Execute Python code',
            inputSchema: { json: { type: 'object', properties: {} } },
          },
        },
        {
          toolSpec: {
            name: 'file_search',
            description: 'Search files',
            inputSchema: { json: { type: 'object', properties: {} } },
          },
        },
      ];

      // Query 1
      const params1 = model.invocationParams({
        tools,
      } satisfies InvocationParamsOptions);

      // Query 2 (same tools)
      const params2 = model.invocationParams({
        tools,
      } satisfies InvocationParamsOptions);

      // Both should have identical tool configs with cachePoint
      expect(params1.toolConfig?.tools).toEqual(params2.toolConfig?.tools);
      expect(params1.toolConfig?.tools).toHaveLength(3); // 2 tools + cachePoint
    });

    it('should handle dynamic tool selection scenario', async () => {
      // Simulate: Query 1 with [A, B] → Query 2 with [A, B, C]
      // Tool set changed, so cache won't match (expected behavior)

      const model = new CustomChatBedrockConverse({
        model: 'anthropic.claude-3-haiku-20240307-v1:0',
        region: 'us-east-1',
        promptCache: true,
      });

      const toolA = {
        toolSpec: {
          name: 'tool_a',
          description: 'Tool A',
          inputSchema: { json: { type: 'object', properties: {} } },
        },
      };

      const toolB = {
        toolSpec: {
          name: 'tool_b',
          description: 'Tool B',
          inputSchema: { json: { type: 'object', properties: {} } },
        },
      };

      const toolC = {
        toolSpec: {
          name: 'tool_c',
          description: 'Tool C',
          inputSchema: { json: { type: 'object', properties: {} } },
        },
      };

      // Query 1: [A, B]
      const params1 = model.invocationParams({
        tools: [toolA, toolB],
      } satisfies InvocationParamsOptions);

      // Query 2: [A, B, C] - different tool set
      const params2 = model.invocationParams({
        tools: [toolA, toolB, toolC],
      } satisfies InvocationParamsOptions);

      // Different lengths (cache prefix won't match beyond common prefix)
      expect(params1.toolConfig?.tools).toHaveLength(3); // 2 tools + cachePoint
      expect(params2.toolConfig?.tools).toHaveLength(4); // 3 tools + cachePoint

      // But both have cachePoint at the end
      expect(params1.toolConfig?.tools?.[2]).toEqual({
        cachePoint: { type: 'default' },
      });
      expect(params2.toolConfig?.tools?.[3]).toEqual({
        cachePoint: { type: 'default' },
      });
    });
  });
});