// ABOUTME: Asserts that countTokens ships the SAME cache_control shape as the
// live request path (_createRequestPayload). This parity test guards against regression.

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { AnthropicProvider } from '../anthropic-provider';
import { Tool } from '@lace/agent/tools/tool';
import { z } from 'zod';
import type { ToolContext, ToolResult } from '@lace/agent/tools/types';
import { anthropicBaseMessagesTrap } from '@lace/agent/test-utils/anthropic-base-namespace-trap';

const mockCreate = vi.fn();
const mockCountTokens = vi.fn();

vi.mock('@anthropic-ai/sdk', () => ({
  default: class MockAnthropic {
    // Throwing trap on base namespace — provider must call beta.messages.*
    messages = anthropicBaseMessagesTrap();
    beta = { messages: { create: mockCreate, stream: vi.fn(), countTokens: mockCountTokens } };
  },
}));

vi.mock('../../utils/logger.js', () => ({
  logger: {
    debug: vi.fn(),
    error: vi.fn(),
    info: vi.fn(),
    trace: vi.fn(),
    warn: vi.fn(),
  },
}));

vi.mock('../../utils/provider-logging.js', () => ({
  logProviderRequest: vi.fn(),
  logProviderResponse: vi.fn(),
}));

class TestTool extends Tool {
  name = 'tool';
  description = 'A tool';
  schema = z.object({ x: z.string() });
  protected async executeValidated(args: { x: string }, _c: ToolContext): Promise<ToolResult> {
    return Promise.resolve(this.createResult(args.x));
  }
}

describe('countTokens shape parity with _createRequestPayload', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mockCountTokens.mockResolvedValue({ input_tokens: 100 });
    mockCreate.mockResolvedValue({
      content: [{ type: 'text', text: 'ok' }],
      usage: { input_tokens: 1, output_tokens: 1 },
    });
  });

  it('countTokens and create receive the same cache_control structure', async () => {
    const provider = new AnthropicProvider({ apiKey: 'k' });
    provider.setSystemPrompt('Sys prompt');

    const messages = [{ role: 'user' as const, content: 'hi' }];
    const tools = [new TestTool()];
    const model = 'claude-sonnet-4-20250514';

    await provider.countTokens(messages, tools, model);
    await provider.createResponse(messages, tools, model);

    expect(mockCountTokens).toHaveBeenCalledTimes(1);
    expect(mockCreate).toHaveBeenCalledTimes(1);

    const countArgs = mockCountTokens.mock.calls[0][0] as {
      system: unknown;
      tools: unknown;
      messages: unknown;
    };
    const createArgs = mockCreate.mock.calls[0][0] as {
      system: unknown;
      tools: unknown;
      messages: unknown;
    };

    // Both must wrap the system prompt in a TextBlockParam array with cache_control.
    expect(countArgs.system).toEqual(createArgs.system);

    // Both must stamp the last (only) tool with cache_control.
    expect(countArgs.tools).toEqual(createArgs.tools);

    // Both must have identical messages.
    expect(countArgs.messages).toEqual(createArgs.messages);
  });
});
