// src/agents/AgentContext.test.ts
import { AgentContext } from './AgentContext';
import { Providers } from '@/common';
import type * as t from '@/types';

describe('AgentContext', () => {
  describe('structured output', () => {
    const baseConfig: t.AgentInputs = {
      agentId: 'test-agent',
      provider: Providers.OPENAI,
      clientOptions: {
        model: 'gpt-4',
        streaming: true,
      },
    };

    describe('isStructuredOutputMode', () => {
      it('should return false when structuredOutput is not configured', () => {
        const context = AgentContext.fromConfig(baseConfig);
        expect(context.isStructuredOutputMode).toBe(false);
      });

      it('should return false when structuredOutput has no schema', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: undefined as unknown as Record<string, unknown>,
          },
        };
        const context = AgentContext.fromConfig(config);
        expect(context.isStructuredOutputMode).toBe(false);
      });

      it('should return false when structuredOutput schema is null', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: null as unknown as Record<string, unknown>,
          },
        };
        const context = AgentContext.fromConfig(config);
        expect(context.isStructuredOutputMode).toBe(false);
      });

      it('should return true when structuredOutput has a valid schema', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: {
              type: 'object',
              properties: {
                answer: { type: 'string' },
              },
            },
          },
        };
        const context = AgentContext.fromConfig(config);
        expect(context.isStructuredOutputMode).toBe(true);
      });

      it('should return true with minimal schema', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: { type: 'string' },
          },
        };
        const context = AgentContext.fromConfig(config);
        expect(context.isStructuredOutputMode).toBe(true);
      });

      // Snake_case structured_output tests (database/API format)
      it('should return true when structured_output (snake_case) is enabled with schema', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structured_output: {
            enabled: true,
            schema: {
              type: 'object',
              properties: {
                result: { type: 'string' },
              },
            },
          },
        };
        const context = AgentContext.fromConfig(config);
        expect(context.isStructuredOutputMode).toBe(true);
      });

      it('should return false when structured_output is not enabled', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structured_output: {
            enabled: false,
            schema: {
              type: 'object',
              properties: {
                result: { type: 'string' },
              },
            },
          },
        };
        const context = AgentContext.fromConfig(config);
        expect(context.isStructuredOutputMode).toBe(false);
      });

      it('should return false when structured_output has no schema', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structured_output: {
            enabled: true,
            // No schema
          },
        };
        const context = AgentContext.fromConfig(config);
        expect(context.isStructuredOutputMode).toBe(false);
      });

      it('should prefer structuredOutput (camelCase) over structured_output (snake_case)', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: {
              type: 'object',
              properties: { camel: { type: 'string' } },
            },
            name: 'CamelSchema',
          },
          structured_output: {
            enabled: true,
            schema: {
              type: 'object',
              properties: { snake: { type: 'string' } },
            },
            name: 'SnakeSchema',
          },
        };
        const context = AgentContext.fromConfig(config);
        expect(context.isStructuredOutputMode).toBe(true);
        expect(context.structuredOutput?.name).toBe('CamelSchema');
      });

      it('should convert structured_output properties correctly', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structured_output: {
            enabled: true,
            schema: {
              type: 'object',
              properties: { data: { type: 'string' } },
            },
            name: 'TestResponse',
            description: 'A test response',
            mode: 'provider',
            strict: true,
          },
        };
        const context = AgentContext.fromConfig(config);
        expect(context.structuredOutput).toEqual({
          schema: {
            type: 'object',
            properties: { data: { type: 'string' } },
          },
          name: 'TestResponse',
          description: 'A test response',
          mode: 'provider',
          strict: true,
        });
      });
    });

    describe('getStructuredOutputSchema', () => {
      it('should return undefined when structuredOutput is not configured', () => {
        const context = AgentContext.fromConfig(baseConfig);
        expect(context.getStructuredOutputSchema()).toBeUndefined();
      });

      it('should return undefined when schema is not set', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {} as t.StructuredOutputConfig,
        };
        const context = AgentContext.fromConfig(config);
        expect(context.getStructuredOutputSchema()).toBeUndefined();
      });

      it('should return the schema as-is when properly configured', () => {
        const schema = {
          type: 'object',
          properties: {
            name: { type: 'string' },
            age: { type: 'number' },
          },
        };
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: { schema },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result).toMatchObject(schema);
      });

      it('should add type: object when properties exist but type is missing', () => {
        const schema = {
          properties: {
            name: { type: 'string' },
          },
        };
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: { schema },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.type).toBe('object');
      });

      it('should add title from config name', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: { type: 'object' },
            name: 'ResponseSchema',
          },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.title).toBe('ResponseSchema');
      });

      it('should not override existing title', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: { type: 'object', title: 'ExistingTitle' },
            name: 'NewName',
          },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.title).toBe('ExistingTitle');
      });

      it('should add description from config', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: { type: 'object' },
            description: 'A structured response',
          },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.description).toBe('A structured response');
      });

      it('should not override existing description', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: { type: 'object', description: 'Original description' },
            description: 'New description',
          },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.description).toBe('Original description');
      });

      it('should set additionalProperties: false in strict mode by default', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: { type: 'object', properties: {} },
          },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.additionalProperties).toBe(false);
      });

      it('should set additionalProperties: false when strict is explicitly true', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: { type: 'object', properties: {} },
            strict: true,
          },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.additionalProperties).toBe(false);
      });

      it('should not set additionalProperties when strict is false', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: { type: 'object', properties: {} },
            strict: false,
          },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.additionalProperties).toBeUndefined();
      });

      it('should preserve existing additionalProperties value', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: {
              type: 'object',
              properties: {},
              additionalProperties: true,
            },
            strict: true,
          },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.additionalProperties).toBe(true);
      });

      it('should not affect non-object types', () => {
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: { type: 'array', items: { type: 'string' } },
          },
        };
        const context = AgentContext.fromConfig(config);
        const result = context.getStructuredOutputSchema();
        expect(result?.additionalProperties).toBeUndefined();
        expect(result?.type).toBe('array');
      });

      it('should not mutate the original schema', () => {
        const originalSchema = {
          type: 'object',
          properties: { name: { type: 'string' } },
        };
        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput: {
            schema: originalSchema,
            name: 'TestSchema',
            description: 'Test description',
          },
        };
        const context = AgentContext.fromConfig(config);
        context.getStructuredOutputSchema();

        // Original should not have been modified
        expect(originalSchema).not.toHaveProperty('title');
        expect(originalSchema).not.toHaveProperty('description');
        expect(originalSchema).not.toHaveProperty('additionalProperties');
      });
    });

    describe('fromConfig with structuredOutput', () => {
      it('should properly initialize structuredOutput from config', () => {
        const structuredOutput: t.StructuredOutputConfig = {
          schema: {
            type: 'object',
            properties: {
              response: { type: 'string' },
              confidence: { type: 'number' },
            },
            required: ['response'],
          },
          name: 'AnalysisResult',
          description: 'The analysis result',
          mode: 'auto',
          strict: true,
        };

        const config: t.AgentInputs = {
          ...baseConfig,
          structuredOutput,
        };

        const context = AgentContext.fromConfig(config);

        expect(context.isStructuredOutputMode).toBe(true);
        expect(context.structuredOutput).toEqual(structuredOutput);
      });

      it('should handle all mode options', () => {
        const modes: Array<t.StructuredOutputMode> = [
          'auto',
          'tool',
          'provider',
        ];

        for (const mode of modes) {
          const config: t.AgentInputs = {
            ...baseConfig,
            structuredOutput: {
              schema: { type: 'object' },
              mode,
            },
          };
          const context = AgentContext.fromConfig(config);
          expect(context.structuredOutput?.mode).toBe(mode);
        }
      });
    });
  });

  describe('summarizeCallback', () => {
    const baseConfig: t.AgentInputs = {
      agentId: 'test-agent',
      provider: Providers.OPENAI,
      clientOptions: {
        model: 'gpt-4',
        streaming: true,
      },
    };

    it('should not have summarizeCallback by default', () => {
      const context = AgentContext.fromConfig(baseConfig);
      expect(context.summarizeCallback).toBeUndefined();
    });

    it('should set summarizeCallback when provided in config', () => {
      const config: t.AgentInputs = {
        ...baseConfig,
        summarizeCallback: async () => 'summary',
      };
      const context = AgentContext.fromConfig(config);
      expect(context.summarizeCallback).toBeDefined();
    });

    it('should preserve summarizeCallback function reference', () => {
      const mockCallback = jest.fn(async () => 'mock summary');
      const config: t.AgentInputs = {
        ...baseConfig,
        summarizeCallback: mockCallback,
      };
      const context = AgentContext.fromConfig(config);
      expect(context.summarizeCallback).toBe(mockCallback);
    });

    it('should allow summarizeCallback to be called', async () => {
      const mockCallback = jest.fn(async () => 'summarized content');
      const config: t.AgentInputs = {
        ...baseConfig,
        summarizeCallback: mockCallback,
      };
      const context = AgentContext.fromConfig(config);
      const result = await context.summarizeCallback!([]);
      expect(result).toBe('summarized content');
      expect(mockCallback).toHaveBeenCalledTimes(1);
    });
  });

  describe('system_cache_blocks (tiered prompt cache)', () => {
    /**
     * `buildSystemRunnable` is private, so we exercise it via the public
     * `initializeSystemRunnable` method and inspect the SystemMessage that
     * gets prepended by the resulting runnable.
     */
    const extractSystemMessage = async (context: AgentContext) => {
      context.initializeSystemRunnable();
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      const runnable = (context as unknown as { cachedSystemRunnable: any })
        .cachedSystemRunnable;
      if (!runnable) return undefined;
      const result = await runnable.invoke([]);
      return result[0];
    };

    it('stores systemCacheBlocks and instructionsCacheTtl on context', () => {
      const config: t.AgentInputs = {
        agentId: 'tier-test',
        provider: Providers.BEDROCK,
        clientOptions: { model: 'anthropic.claude-3-5-sonnet-20241022-v2:0' },
        instructions: 'agent specific',
        system_cache_blocks: [
          { text: 'platform tier text', ttl: '5m' },
        ],
        instructions_cache_ttl: '1h',
      };
      const context = AgentContext.fromConfig(config);
      expect(context.systemCacheBlocks).toEqual([
        { text: 'platform tier text', ttl: '5m' },
      ]);
      expect(context.instructionsCacheTtl).toBe('1h');
    });

    it('emits multi-block SystemMessage with single cachePoint after Tier 1 for Bedrock Claude', async () => {
      const config: t.AgentInputs = {
        agentId: 'bedrock-claude',
        provider: Providers.BEDROCK,
        clientOptions: {
          model: 'anthropic.claude-3-5-sonnet-20241022-v2:0',
          promptCache: true,
        } as t.BedrockAnthropicClientOptions,
        instructions: 'AGENT_TIER_2',
        system_cache_blocks: [{ text: 'PLATFORM_TIER_1' }],
      };
      const context = AgentContext.fromConfig(config);
      const systemMessage = await extractSystemMessage(context);
      expect(systemMessage).toBeDefined();
      const content = systemMessage?.content as Array<Record<string, unknown>>;
      expect(Array.isArray(content)).toBe(true);
      // Only the platform tier gets an explicit cachePoint; the trailing
      // per-agent block is uncached at the system layer (tools breakpoint
      // covers it via forward-prefix-hash).
      expect(content).toEqual([
        { type: 'text', text: 'PLATFORM_TIER_1' },
        { cachePoint: { type: 'default' } },
        { type: 'text', text: expect.stringContaining('AGENT_TIER_2') },
      ]);
    });

    it('emits cache_control on Tier 1 text block only for Anthropic with promptCache=true', async () => {
      const config: t.AgentInputs = {
        agentId: 'anthropic-direct',
        provider: Providers.ANTHROPIC,
        clientOptions: {
          model: 'claude-3-5-sonnet-20241022',
          promptCache: true,
        } as t.AnthropicClientOptions,
        instructions: 'AGENT_TIER_2',
        system_cache_blocks: [{ text: 'PLATFORM_TIER_1' }],
      };
      const context = AgentContext.fromConfig(config);
      const systemMessage = await extractSystemMessage(context);
      const content = systemMessage?.content as Array<Record<string, unknown>>;
      expect(content).toEqual([
        {
          type: 'text',
          text: 'PLATFORM_TIER_1',
          cache_control: { type: 'ephemeral' },
        },
        {
          type: 'text',
          text: expect.stringContaining('AGENT_TIER_2'),
        },
      ]);
    });

    it('falls back to single string SystemMessage for Bedrock Nova (no cachePoint support)', async () => {
      const config: t.AgentInputs = {
        agentId: 'bedrock-nova',
        provider: Providers.BEDROCK,
        clientOptions: {
          model: 'amazon.nova-pro-v1:0',
          promptCache: true,
        } as unknown as t.BedrockAnthropicClientOptions,
        instructions: 'AGENT_TIER_2',
        system_cache_blocks: [{ text: 'PLATFORM_TIER_1' }],
      };
      const context = AgentContext.fromConfig(config);
      const systemMessage = await extractSystemMessage(context);
      // Non-Claude Bedrock: tiered emission disabled, content stays string
      expect(typeof systemMessage?.content).toBe('string');
      expect(systemMessage?.content).toContain('AGENT_TIER_2');
      // platform tier is NOT in the system message because the lib can't
      // emit a cachePoint for it on Nova; ranger's quick-patch puts the
      // bytes into `instructions` for Nova-style providers.
      expect(systemMessage?.content).not.toContain('PLATFORM_TIER_1');
    });

    it('falls back to single string for Bedrock Claude when promptCache is false', async () => {
      const config: t.AgentInputs = {
        agentId: 'bedrock-claude-nopc',
        provider: Providers.BEDROCK,
        clientOptions: {
          model: 'anthropic.claude-3-5-sonnet-20241022-v2:0',
          promptCache: false,
        } as t.BedrockAnthropicClientOptions,
        instructions: 'AGENT_TIER_2',
        system_cache_blocks: [{ text: 'PLATFORM_TIER_1' }],
      };
      const context = AgentContext.fromConfig(config);
      const systemMessage = await extractSystemMessage(context);
      expect(typeof systemMessage?.content).toBe('string');
    });

    it('preserves legacy single-block Anthropic caching when no system_cache_blocks set', async () => {
      const config: t.AgentInputs = {
        agentId: 'anthropic-legacy',
        provider: Providers.ANTHROPIC,
        clientOptions: {
          model: 'claude-3-5-sonnet-20241022',
          promptCache: true,
        } as t.AnthropicClientOptions,
        instructions: 'AGENT_INSTR',
      };
      const context = AgentContext.fromConfig(config);
      const systemMessage = await extractSystemMessage(context);
      const content = systemMessage?.content as Array<Record<string, unknown>>;
      expect(content).toEqual([
        {
          type: 'text',
          text: expect.stringContaining('AGENT_INSTR'),
          cache_control: { type: 'ephemeral' },
        },
      ]);
    });

    it('emits multiple Tier 1 text blocks but only ONE cachePoint at the boundary', async () => {
      const config: t.AgentInputs = {
        agentId: 'multi-tier',
        provider: Providers.BEDROCK,
        clientOptions: {
          model: 'anthropic.claude-3-5-sonnet-20241022-v2:0',
          promptCache: true,
        } as t.BedrockAnthropicClientOptions,
        instructions: 'PER_AGENT',
        system_cache_blocks: [
          { text: 'TIER1_PLATFORM' },
          { text: 'TIER1_CAPABILITIES' },
        ],
      };
      const context = AgentContext.fromConfig(config);
      const systemMessage = await extractSystemMessage(context);
      const content = systemMessage?.content as Array<Record<string, unknown>>;
      // Multiple Tier 1 entries collapse into the same cached prefix —
      // only one cachePoint is emitted (at the Tier1→Tier2 boundary) so
      // we don't blow the 4-breakpoint cap from Anthropic/Bedrock.
      expect(content).toEqual([
        { type: 'text', text: 'TIER1_PLATFORM' },
        { type: 'text', text: 'TIER1_CAPABILITIES' },
        { cachePoint: { type: 'default' } },
        { type: 'text', text: expect.stringContaining('PER_AGENT') },
      ]);
    });
  });
});