import { simulateReadableStream } from "ai";
import { MockLanguageModelV3 } from "ai/test";

const MOCK_USAGE = {
  inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined },
  outputTokens: { total: 5, text: 5, reasoning: undefined },
} as const;

const STOP_FINISH = {
  type: "finish" as const,
  finishReason: { unified: "stop" as const, raw: undefined },
  usage: MOCK_USAGE,
};

const TOOL_CALLS_FINISH = {
  type: "finish" as const,
  finishReason: { unified: "tool-calls" as const, raw: undefined },
  usage: MOCK_USAGE,
};

type StreamChunk = Parameters<typeof simulateReadableStream>[0]["chunks"][number];
type StreamResult = { stream: ReadableStream };

export function textStreamResult(content: string): () => StreamResult {
  return () => {
    const textId = `text-${crypto.randomUUID().slice(0, 8)}`;
    const chunks: StreamChunk[] = [
      { type: "text-start", id: textId },
      { type: "text-delta", id: textId, delta: content },
      { type: "text-end", id: textId },
      STOP_FINISH,
    ];
    return {
      stream: simulateReadableStream({ chunks, initialDelayInMs: null, chunkDelayInMs: null }),
    };
  };
}

export function toolCallStreamResult(
  calls: Array<{ id?: string; name: string; arguments: Record<string, unknown> }>,
): () => StreamResult {
  return () => {
    const chunks: StreamChunk[] = [];
    for (const c of calls) {
      const toolCallId = c.id ?? `call_${crypto.randomUUID().slice(0, 8)}`;
      chunks.push({
        type: "tool-call",
        toolCallId,
        toolName: c.name,
        input: JSON.stringify(c.arguments),
      });
    }
    chunks.push(TOOL_CALLS_FINISH);
    return {
      stream: simulateReadableStream({ chunks, initialDelayInMs: null, chunkDelayInMs: null }),
    };
  };
}

export function textGenerateResult(content: string) {
  return {
    content: [{ type: "text" as const, text: content }],
    finishReason: { unified: "stop" as const, raw: undefined },
    usage: MOCK_USAGE,
    warnings: [] as never[],
  };
}

/**
 * Create a MockLanguageModelV3 with queued responses.
 *
 * Uses function-based doStream/doGenerate to work around the
 * off-by-one bug in MockLanguageModelV3's array indexing.
 */
export function createMockModel(opts?: {
  streamResponses?: Array<() => StreamResult>;
  generateResponses?: ReturnType<typeof textGenerateResult>[];
}) {
  let streamIdx = 0;
  let genIdx = 0;
  const streamQueue = opts?.streamResponses ?? [];
  const genQueue = opts?.generateResponses ?? [];

  return new MockLanguageModelV3({
    doStream:
      streamQueue.length > 0
        ? async () => {
            const factory = streamQueue[streamIdx];
            if (!factory) throw new Error(`MockModel: no more queued stream responses (called ${streamIdx + 1} times)`);
            streamIdx++;
            return factory();
          }
        : undefined,
    doGenerate:
      genQueue.length > 0
        ? async () => {
            const result = genQueue[genIdx];
            if (!result) throw new Error(`MockModel: no more queued generate responses (called ${genIdx + 1} times)`);
            genIdx++;
            return result;
          }
        : undefined,
  });
}