path: root/packages/cli/src/core/gemini-client.ts

import {
  GenerateContentConfig,
  GoogleGenAI,
  Part,
  Chat,
  Type,
  SchemaUnion,
  PartListUnion,
  Content,
} from '@google/genai';
import { getApiKey } from '../config/env.js';
import { CoreSystemPrompt } from './prompts.js';
import {
  type ToolCallEvent,
  type ToolCallConfirmationDetails,
  ToolCallStatus,
} from '../ui/types.js';
import process from 'node:process';
import { toolRegistry } from '../tools/tool-registry.js';
import { ToolResult } from '../tools/tools.js';
import { getFolderStructure } from '../utils/getFolderStructure.js';
import { GeminiEventType, GeminiStream } from './gemini-stream.js';

type ToolExecutionOutcome = {
  callId: string;
  name: string;
  args: Record<string, any>;
  result?: ToolResult;
  error?: any;
  confirmationDetails?: ToolCallConfirmationDetails;
};

export class GeminiClient {
  private ai: GoogleGenAI;
  private defaultHyperParameters: GenerateContentConfig = {
    temperature: 0,
    topP: 1,
  };
  private readonly MAX_TURNS = 100;

  constructor() {
    const apiKey = getApiKey();
    this.ai = new GoogleGenAI({ apiKey });
  }

  public async startChat(): Promise<Chat> {
    const tools = toolRegistry.getToolSchemas();

    // --- Get environmental information ---
    const cwd = process.cwd();
    const today = new Date().toLocaleDateString(undefined, {
      // Use locale-aware date formatting
      weekday: 'long',
      year: 'numeric',
      month: 'long',
      day: 'numeric',
    });
    const platform = process.platform;

    // --- Format information into a conversational multi-line string ---
    const folderStructure = await getFolderStructure(cwd);
    // --- End folder structure formatting ---)
    const initialContextText = `
Okay, just setting up the context for our chat.
Today is ${today}.
My operating system is: ${platform}
I'm currently working in the directory: ${cwd}
${folderStructure}
        `.trim();

    const initialContextPart: Part = { text: initialContextText };
    // --- End environmental information formatting ---

    try {
      const chat = this.ai.chats.create({
        model: 'gemini-2.0-flash', //'gemini-2.0-flash',
        config: {
          systemInstruction: CoreSystemPrompt,
          ...this.defaultHyperParameters,
          tools,
        },
        history: [
          // --- Add the context as a single part in the initial user message ---
          {
            role: 'user',
            parts: [initialContextPart], // Pass the single Part object in an array
          },
          // --- Add an empty model response to balance the history ---
          {
            role: 'model',
            parts: [{ text: 'Got it. Thanks for the context!' }], // A slightly more conversational model response
          },
          // --- End history modification ---
        ],
      });
      return chat;
    } catch (error) {
      console.error('Error initializing Gemini chat session:', error);
      const message = error instanceof Error ? error.message : 'Unknown error.';
      throw new Error(`Failed to initialize chat: ${message}`);
    }
  }

  public addMessageToHistory(chat: Chat, message: Content): void {
    const history = chat.getHistory();
    history.push(message);
    this.ai.chats;
    chat;
  }

  public async *sendMessageStream(
    chat: Chat,
    request: PartListUnion,
    signal?: AbortSignal,
  ): GeminiStream {
    let currentMessageToSend: PartListUnion = request;
    let turns = 0;

    try {
      while (turns < this.MAX_TURNS) {
        turns++;
        const resultStream = await chat.sendMessageStream({
          message: currentMessageToSend,
        });
        let functionResponseParts: Part[] = [];
        let pendingToolCalls: Array<{
          callId: string;
          name: string;
          args: Record<string, any>;
        }> = [];
        let yieldedTextInTurn = false;
        const chunksForDebug = [];

        for await (const chunk of resultStream) {
          chunksForDebug.push(chunk);
          if (signal?.aborted) {
            const abortError = new Error(
              'Request cancelled by user during stream.',
            );
            abortError.name = 'AbortError';
            throw abortError;
          }

          const functionCalls = chunk.functionCalls;
          if (functionCalls && functionCalls.length > 0) {
            for (const call of functionCalls) {
              const callId =
                call.id ??
                `${call.name}-${Date.now()}-${Math.random().toString(16).slice(2)}`;
              const name = call.name || 'undefined_tool_name';
              const args = (call.args || {}) as Record<string, any>;

              pendingToolCalls.push({ callId, name, args });
              const evtValue: ToolCallEvent = {
                type: 'tool_call',
                status: ToolCallStatus.Pending,
                callId,
                name,
                args,
                resultDisplay: undefined,
                confirmationDetails: undefined,
              };
              yield {
                type: GeminiEventType.ToolCallInfo,
                value: evtValue,
              };
            }
          } else {
            const text = chunk.text;
            if (text) {
              yieldedTextInTurn = true;
              yield {
                type: GeminiEventType.Content,
                value: text,
              };
            }
          }
        }

        if (pendingToolCalls.length > 0) {
          const toolPromises: Promise<ToolExecutionOutcome>[] =
            pendingToolCalls.map(async (pendingToolCall) => {
              const tool = toolRegistry.getTool(pendingToolCall.name);

              if (!tool) {
                // Directly return error outcome if tool not found
                return {
                  ...pendingToolCall,
                  error: new Error(
                    `Tool "${pendingToolCall.name}" not found or is not registered.`,
                  ),
                };
              }

              try {
                const confirmation = await tool.shouldConfirmExecute(
                  pendingToolCall.args,
                );
                if (confirmation) {
                  return {
                    ...pendingToolCall,
                    confirmationDetails: confirmation,
                  };
                }
              } catch (error) {
                return {
                  ...pendingToolCall,
                  error: new Error(
                    `Tool failed to check tool confirmation: ${error}`,
                  ),
                };
              }

              try {
                const result = await tool.execute(pendingToolCall.args);
                return { ...pendingToolCall, result };
              } catch (error) {
                return {
                  ...pendingToolCall,
                  error: new Error(`Tool failed to execute: ${error}`),
                };
              }
            });
          const toolExecutionOutcomes: ToolExecutionOutcome[] =
            await Promise.all(toolPromises);

          for (const executedTool of toolExecutionOutcomes) {
            const { callId, name, args, result, error, confirmationDetails } =
              executedTool;

            if (error) {
              const errorMessage = error?.message || String(error);
              yield {
                type: GeminiEventType.Content,
                value: `[Error invoking tool ${name}: ${errorMessage}]`,
              };
            } else if (
              result &&
              typeof result === 'object' &&
              result !== null &&
              'error' in result
            ) {
              const errorMessage = String(result.error);
              yield {
                type: GeminiEventType.Content,
                value: `[Error executing tool ${name}: ${errorMessage}]`,
              };
            } else {
              const status = confirmationDetails
                ? ToolCallStatus.Confirming
                : ToolCallStatus.Invoked;
              const evtValue: ToolCallEvent = {
                type: 'tool_call',
                status,
                callId,
                name,
                args,
                resultDisplay: result?.returnDisplay,
                confirmationDetails,
              };
              yield {
                type: GeminiEventType.ToolCallInfo,
                value: evtValue,
              };
            }
          }

          pendingToolCalls = [];

          const waitingOnConfirmations =
            toolExecutionOutcomes.filter(
              (outcome) => outcome.confirmationDetails,
            ).length > 0;
          if (waitingOnConfirmations) {
            // Stop processing content, wait for user.
            // TODO: Kill token processing once API supports signals.
            break;
          }

          functionResponseParts = toolExecutionOutcomes.map(
            (executedTool: ToolExecutionOutcome): Part => {
              const { name, result, error } = executedTool;
              const output = { output: result?.llmContent };
              let toolOutcomePayload: any;

              if (error) {
                const errorMessage = error?.message || String(error);
                toolOutcomePayload = {
                  error: `Invocation failed: ${errorMessage}`,
                };
                console.error(
                  `[Turn ${turns}] Critical error invoking tool ${name}:`,
                  error,
                );
              } else if (
                result &&
                typeof result === 'object' &&
                result !== null &&
                'error' in result
              ) {
                toolOutcomePayload = output;
                console.warn(
                  `[Turn ${turns}] Tool ${name} returned an error structure:`,
                  result.error,
                );
              } else {
                toolOutcomePayload = output;
              }

              return {
                functionResponse: {
                  name: name,
                  id: executedTool.callId,
                  response: toolOutcomePayload,
                },
              };
            },
          );
          currentMessageToSend = functionResponseParts;
        } else if (yieldedTextInTurn) {
          const history = chat.getHistory();
          const checkPrompt = `Analyze *only* the content and structure of your immediately preceding response (your last turn in the conversation history). Based *strictly* on that response, determine who should logically speak next: the 'user' or the 'model' (you).

**Decision Rules (apply in order):**

1.  **Model Continues:** If your last response explicitly states an immediate next action *you* intend to take (e.g., "Next, I will...", "Now I'll process...", "Moving on to analyze...", indicates an intended tool call that didn't execute), OR if the response seems clearly incomplete (cut off mid-thought without a natural conclusion), then the **'model'** should speak next.
2.  **Question to User:** If your last response ends with a direct question specifically addressed *to the user*, then the **'user'** should speak next.
3.  **Waiting for User:** If your last response completed a thought, statement, or task *and* does not meet the criteria for Rule 1 (Model Continues) or Rule 2 (Question to User), it implies a pause expecting user input or reaction. In this case, the **'user'** should speak next.

**Output Format:**

Respond *only* in JSON format according to the following schema. Do not include any text outside the JSON structure.

\`\`\`json
{
  "type": "object",
  "properties": {
    "reasoning": {
        "type": "string",
        "description": "Brief explanation justifying the 'next_speaker' choice based *strictly* on the applicable rule and the content/structure of the preceding turn."
    },
    "next_speaker": {
      "type": "string",
      "enum": ["user", "model"],
      "description": "Who should speak next based *only* on the preceding turn and the decision rules."
    }
  },
  "required": ["next_speaker", "reasoning"]
\`\`\`
}`;

          // Schema Idea
          const responseSchema: SchemaUnion = {
            type: Type.OBJECT,
            properties: {
              reasoning: {
                type: Type.STRING,
                description:
                  "Brief explanation justifying the 'next_speaker' choice based *strictly* on the applicable rule and the content/structure of the preceding turn.",
              },
              next_speaker: {
                type: Type.STRING,
                enum: ['user', 'model'], // Enforce the choices
                description:
                  'Who should speak next based *only* on the preceding turn and the decision rules',
              },
            },
            required: ['reasoning', 'next_speaker'],
          };

          try {
            // Use the new generateJson method, passing the history and the check prompt
            const parsedResponse = await this.generateJson(
              [
                ...history,
                {
                  role: 'user',
                  parts: [{ text: checkPrompt }],
                },
              ],
              responseSchema,
            );

            // Safely extract the next speaker value
            const nextSpeaker: string | undefined =
              typeof parsedResponse?.next_speaker === 'string'
                ? parsedResponse.next_speaker
                : undefined;

            if (nextSpeaker === 'model') {
              currentMessageToSend = { text: 'alright' }; // Or potentially a more meaningful continuation prompt
            } else {
              // 'user' should speak next, or value is missing/invalid. End the turn.
              break;
            }
          } catch (error) {
            console.error(
              `[Turn ${turns}] Failed to get or parse next speaker check:`,
              error,
            );
            // If the check fails, assume user should speak next to avoid infinite loops
            break;
          }
        } else {
          console.warn(
            `[Turn ${turns}] No text or function calls received from Gemini. Ending interaction.`,
          );
          break;
        }
      }

      if (turns >= this.MAX_TURNS) {
        console.warn(
          'sendMessageStream: Reached maximum tool call turns limit.',
        );
        yield {
          type: GeminiEventType.Content,
          value:
            '\n\n[System Notice: Maximum interaction turns reached. The conversation may be incomplete.]',
        };
      }
    } catch (error: unknown) {
      if (error instanceof Error && error.name === 'AbortError') {
        console.log('Gemini stream request aborted by user.');
        throw error;
      } else {
        console.error(`Error during Gemini stream or tool interaction:`, error);
        const message = error instanceof Error ? error.message : String(error);
        yield {
          type: GeminiEventType.Content,
          value: `\n\n[Error: An unexpected error occurred during the chat: ${message}]`,
        };
        throw error;
      }
    }
  }

  /**
   * Generates structured JSON content based on conversational history and a schema.
   * @param contents The conversational history (Content array) to provide context.
   * @param schema The SchemaUnion defining the desired JSON structure.
   * @returns A promise that resolves to the parsed JSON object matching the schema.
   * @throws Throws an error if the API call fails or the response is not valid JSON.
   */
  public async generateJson(
    contents: Content[],
    schema: SchemaUnion,
  ): Promise<any> {
    try {
      const result = await this.ai.models.generateContent({
        model: 'gemini-2.0-flash', // Using flash for potentially faster structured output
        config: {
          ...this.defaultHyperParameters,
          systemInstruction: CoreSystemPrompt,
          responseSchema: schema,
          responseMimeType: 'application/json',
        },
        contents: contents, // Pass the full Content array
      });

      const responseText = result.text;
      if (!responseText) {
        throw new Error('API returned an empty response.');
      }

      try {
        const parsedJson = JSON.parse(responseText);
        // TODO: Add schema validation if needed
        return parsedJson;
      } catch (parseError) {
        console.error('Failed to parse JSON response:', responseText);
        throw new Error(
          `Failed to parse API response as JSON: ${parseError instanceof Error ? parseError.message : String(parseError)}`,
        );
      }
    } catch (error) {
      console.error('Error generating JSON content:', error);
      const message =
        error instanceof Error ? error.message : 'Unknown API error.';
      throw new Error(`Failed to generate JSON content: ${message}`);
    }
  }
}