/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import {
  GenerateContentResponse,
  type CountTokensResponse,
  type GenerateContentParameters,
  type CountTokensParameters,
  EmbedContentResponse,
  type EmbedContentParameters,
} from '@google/genai';
import { promises } from 'node:fs';
import type { ContentGenerator } from './contentGenerator.js';
import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
import type { LlmRole } from '../telemetry/types.js';

export type FakeResponse =
  | {
      method: 'generateContent';
      response: GenerateContentResponse;
    }
  | {
      method: 'generateContentStream';
      response: GenerateContentResponse[];
    }
  | {
      method: 'countTokens';
      response: CountTokensResponse;
    }
  | {
      method: 'embedContent';
      response: EmbedContentResponse;
    };

/**
 * Options for the FakeContentGenerator.
 */
export interface FakeContentGeneratorOptions {
  /**
   * If true, the generator will find the first available response that matches
   * the requested method, rather than strictly following the input order.
   * Useful for non-deterministic background tasks.
   */
  nonStrict?: boolean;
}

// A ContentGenerator that responds with canned responses.
//
// Typically these would come from a file, provided by the `--fake-responses`
// CLI argument.
export class FakeContentGenerator implements ContentGenerator {
  private callCounter = 0;
  userTier?: UserTierId;
  userTierName?: string;
  paidTier?: GeminiUserTier;

  private readonly responses: FakeResponse[];

  constructor(
    responses: FakeResponse[],
    private readonly options: FakeContentGeneratorOptions = {},
  ) {
    this.responses = structuredClone(responses);
  }

  static async fromFile(
    filePath: string,
    options: FakeContentGeneratorOptions = {},
  ): Promise<FakeContentGenerator> {
    const fileContent = await promises.readFile(filePath, 'utf-8');
    const responses = fileContent
      .split('\n')
      .filter((line) => line.trim() !== '')
      // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
      .map((line) => JSON.parse(line) as FakeResponse);
    return new FakeContentGenerator(responses, options);
  }

  private getNextResponse<
    M extends FakeResponse['method'],
    R = Extract<FakeResponse, { method: M }>['response'],
  >(method: M, request: unknown): R {
    if (this.options.nonStrict) {
      const index = this.responses.findIndex((r) => r.method === method);
      if (index === -1) {
        throw new Error(
          `No more mock responses for ${method}, got request:\n` +
            safeJsonStringify(request),
        );
      }
      const response = this.responses.splice(index, 1)[0];
      // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
      return response.response as R;
    }

    const response = this.responses[this.callCounter++];
    if (!response) {
      throw new Error(
        `No more mock responses for ${method}, got request:\n` +
          safeJsonStringify(request),
      );
    }
    if (response.method !== method) {
      throw new Error(
        `Unexpected response type, next response was for ${response.method} but expected ${method}`,
      );
    }
    // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
    return response.response as R;
  }

  async generateContent(
    request: GenerateContentParameters,
    _userPromptId: string,
    // eslint-disable-next-line @typescript-eslint/no-unused-vars
    role: LlmRole,
  ): Promise<GenerateContentResponse> {
    const response: unknown = this.getNextResponse('generateContent', request);
    Object.setPrototypeOf(response, GenerateContentResponse.prototype);
    if (response instanceof GenerateContentResponse) {
      return response;
    }
    throw new Error('Failed to create GenerateContentResponse');
  }

  async generateContentStream(
    request: GenerateContentParameters,
    _userPromptId: string,
    // eslint-disable-next-line @typescript-eslint/no-unused-vars
    role: LlmRole,
  ): Promise<AsyncGenerator<GenerateContentResponse>> {
    const responses = this.getNextResponse('generateContentStream', request);
    async function* stream() {
      for (const response of responses) {
        yield Object.setPrototypeOf(
          response,
          GenerateContentResponse.prototype,
        );
      }
    }
    return stream();
  }

  async countTokens(
    request: CountTokensParameters,
  ): Promise<CountTokensResponse> {
    return this.getNextResponse('countTokens', request);
  }

  async embedContent(
    request: EmbedContentParameters,
  ): Promise<EmbedContentResponse> {
    const response: unknown = this.getNextResponse('embedContent', request);
    Object.setPrototypeOf(response, EmbedContentResponse.prototype);
    if (response instanceof EmbedContentResponse) {
      return response;
    }
    throw new Error('Failed to create EmbedContentResponse');
  }
}
