import { describe, it, expect, vi } from 'vitest';
import { z } from 'zod';
import { callLLMStage } from './call';
import { HAIKU } from './models';
import type Anthropic from '@anthropic-ai/sdk';

/**
 * Minimal in-memory fake of `client.messages.create`. Each call dequeues the
 * next response from `queue`. Tracks bodies for assertion.
 */
function makeFakeClient(queue: Array<{ text: string; usage?: Partial<Anthropic.Usage> }>) {
  const calls: Array<{ body: any }> = [];
  const create = vi.fn(async (body: any) => {
    calls.push({ body });
    const next = queue.shift();
    if (!next) throw new Error('Fake client: queue exhausted');
    return {
      id: `msg_${calls.length}`,
      type: 'message',
      role: 'assistant',
      content: [{ type: 'text', text: next.text }],
      model: body.model,
      stop_reason: 'end_turn',
      stop_sequence: null,
      usage: {
        input_tokens: next.usage?.input_tokens ?? 100,
        output_tokens: next.usage?.output_tokens ?? 50,
        cache_creation_input_tokens: next.usage?.cache_creation_input_tokens ?? 0,
        cache_read_input_tokens: next.usage?.cache_read_input_tokens ?? 0,
      },
    } as Anthropic.Message;
  });

  const fake = {
    messages: { create },
  } as unknown as Anthropic;

  return { fake, calls, create };
}

const PingSchema = z.object({ ok: z.literal(true), n: z.number() });

describe('callLLMStage', () => {
  it('returns parsed data on first attempt when response is valid JSON', async () => {
    const { fake, create } = makeFakeClient([
      { text: '{"ok": true, "n": 7}' },
    ]);

    const result = await callLLMStage({
      client: fake,
      model: HAIKU,
      system: 'sys',
      user: 'usr',
      schema: PingSchema,
      maxTokens: 1024,
    });

    expect(result.attempts).toBe(1);
    expect(result.data).toEqual({ ok: true, n: 7 });
    expect(result.usage.input_tokens).toBe(100);
    expect(result.estimated_cost_usd).toBeGreaterThan(0);
    expect(create).toHaveBeenCalledTimes(1);
  });

  it('strips a ```json code fence on first attempt without retrying', async () => {
    const { fake, create } = makeFakeClient([
      { text: '```json\n{"ok": true, "n": 1}\n```' },
    ]);
    const result = await callLLMStage({
      client: fake,
      model: HAIKU,
      system: 'sys',
      user: 'usr',
      schema: PingSchema,
      maxTokens: 1024,
    });
    expect(result.attempts).toBe(1);
    expect(result.data).toEqual({ ok: true, n: 1 });
    expect(create).toHaveBeenCalledTimes(1);
  });

  it('retries once when first response fails JSON.parse', async () => {
    const { fake, create, calls } = makeFakeClient([
      { text: 'this is not json' },
      { text: '{"ok": true, "n": 2}' },
    ]);

    const result = await callLLMStage({
      client: fake,
      model: HAIKU,
      system: 'sys',
      user: 'usr',
      schema: PingSchema,
      maxTokens: 1024,
    });

    expect(result.attempts).toBe(2);
    expect(result.data).toEqual({ ok: true, n: 2 });
    expect(create).toHaveBeenCalledTimes(2);

    // The second call's history must include the assistant's bad response and
    // a follow-up user turn explaining the error.
    const secondBody = calls[1].body;
    expect(secondBody.messages).toHaveLength(3);
    expect(secondBody.messages[0].role).toBe('user');
    expect(secondBody.messages[1].role).toBe('assistant');
    expect(secondBody.messages[1].content).toBe('this is not json');
    expect(secondBody.messages[2].role).toBe('user');
    expect(secondBody.messages[2].content).toMatch(/could not be parsed/);
  });

  it('retries once when first response fails Zod validation', async () => {
    const { fake, create } = makeFakeClient([
      { text: '{"ok": true, "n": "should-be-number"}' },
      { text: '{"ok": true, "n": 3}' },
    ]);

    const result = await callLLMStage({
      client: fake,
      model: HAIKU,
      system: 'sys',
      user: 'usr',
      schema: PingSchema,
      maxTokens: 1024,
    });

    expect(result.attempts).toBe(2);
    expect(result.data).toEqual({ ok: true, n: 3 });
    expect(create).toHaveBeenCalledTimes(2);
  });

  it('throws after two failures, including both error messages', async () => {
    const { fake } = makeFakeClient([
      { text: 'not json' },
      { text: 'still not json' },
    ]);

    await expect(
      callLLMStage({
        client: fake,
        model: HAIKU,
        system: 'sys',
        user: 'usr',
        schema: PingSchema,
        maxTokens: 1024,
      })
    ).rejects.toThrow(/failed twice/);
  });

  it('sums usage across attempts on retry', async () => {
    const { fake } = makeFakeClient([
      { text: 'not json', usage: { input_tokens: 100, output_tokens: 10 } },
      { text: '{"ok": true, "n": 4}', usage: { input_tokens: 120, output_tokens: 20 } },
    ]);

    const result = await callLLMStage({
      client: fake,
      model: HAIKU,
      system: 'sys',
      user: 'usr',
      schema: PingSchema,
      maxTokens: 1024,
    });

    expect(result.usage.input_tokens).toBe(220);
    expect(result.usage.output_tokens).toBe(30);
  });

  it('marks the system prompt with cache_control: ephemeral', async () => {
    const { fake, calls } = makeFakeClient([{ text: '{"ok": true, "n": 5}' }]);
    await callLLMStage({
      client: fake,
      model: HAIKU,
      system: 'sys',
      user: 'usr',
      schema: PingSchema,
      maxTokens: 1024,
    });
    expect(calls[0].body.system).toEqual([
      { type: 'text', text: 'sys', cache_control: { type: 'ephemeral' } },
    ]);
  });

  it('does NOT pass temperature/top_p/top_k (Opus 4.7 would 400)', async () => {
    const { fake, calls } = makeFakeClient([{ text: '{"ok": true, "n": 6}' }]);
    await callLLMStage({
      client: fake,
      model: HAIKU,
      system: 'sys',
      user: 'usr',
      schema: PingSchema,
      maxTokens: 1024,
    });
    const body = calls[0].body;
    expect(body.temperature).toBeUndefined();
    expect(body.top_p).toBeUndefined();
    expect(body.top_k).toBeUndefined();
  });
});