// Updated BaseAgent class to also include the ingestResponseAndStream method,
// allowing derived classes to call it, rather than defining their own.
//
// Additionally, we leave the fitHistoryWithinTokenBudget method as is.
//
// We'll remove the local ingestResponseAndStream from PluginAgent in a subsequent step.

import { Response, Request } from "express";
import Anthropic from "@anthropic-ai/sdk";
import path from "path";
import {
  getChat,
  createChat,
  persistChatMessage,
  getChatMessages,
  sequenceToolCallAndResponses,
  AgentProgramType,
  toInMemoryChatMessage,
  ApptileInMemoryChatMessage,
  toAnthropicMessages,
  toOpenaiMessages,
  getAPIKey,
  ApptileChatMessage,
  toAmazonMessages,
} from "../../database/cliconfig";
import { readFile } from "node:fs/promises";
import { nodejsRoot } from "../../projectrootpath";
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import { Logger, logger } from "../../utils/logger";
import { ingestAmazonMessageStream, ingestClaudeResponseStream, ingestOpenaiMessageStream } from "../messageparsing";
import OpenAI from "openai";
import { encoding_for_model, TiktokenModel } from "tiktoken";
import { ChatCompletionMessageParam, ChatCompletionTool } from "openai/resources/chat";
import chalk from "chalk";
import { ITokenCounterRequest, TokenCounter } from "../tokenCounter";
import { SupabaseCreds } from "./mcp-servers/supabase";
import { BedrockRuntimeClient, ConverseStreamCommand, Tool as BedrockTool, ToolInputSchema } from "@aws-sdk/client-bedrock-runtime";
import zodToJsonSchema from "zod-to-json-schema";
import { readAllPluginFilesByPluginName } from "./mcp-servers/tools";
import { v4 } from "uuid";

export const INSERTED_CODE_PREFIX = `This message is added before each call to the agent and contains the code files that exist so far and is not visible in the chat client. This is the current code:`

// Array of colors you want to cycle through
const colors = ['red', 'yellow', 'blue', 'magenta', 'cyan'];
let currentColor = 0;
export function getColor() {
  return colors[currentColor];
}

export function cycleColor() {
  currentColor = (currentColor + 1) % colors.length;
  return chalk.keyword(colors[currentColor]);
}

export type MCPClient = {
  client: Client;
};

export type NextStep = "ContinueChat" | "WaitForUserInput" | "ToolError";

export type LLMProviders = "openai" | "claude" | "google" | "amazon";
export function narrowLLMProvider(provider: string) : provider is LLMProviders {
  return ["openai", "claude", "google", "amazon"].includes(provider);
}

export type GoogleModels = "gemini-2.5-pro" | "gemini-2.5-flash-preview-05-20";

export function narrowGoogleModel(model: string) : model is GoogleModels {
  return [
    "gemini-2.5-pro",
    "gemini-2.5-flash-preview-05-20"
  ].includes(model);
}

export type AnthropicModels = 'claude-sonnet-4-20250514'|
                        'claude-3-7-sonnet-20250219'|
                        'claude-3-5-sonnet-latest'|
                        'claude-3-5-haiku-latest';

export function narrowAnthropicModel(model: any) : model is AnthropicModels {
  return [
    'claude-sonnet-4-20250514', 
    'claude-3-7-sonnet-20250219',
    'claude-3-5-sonnet-latest',
    'claude-3-5-haiku-latest'
  ].includes(model);
}

export type OpenAIModels = 'gpt-4.1'|'o1'|'o1-pro'|'gpt-4o'|'o4-mini';

export function narrowOpenAIModel(model: any) : model is OpenAIModels {
  return ['gpt-4.1', 'o1', 'o1-pro', 'gpt-4o', 'o4-mini'].includes(model);
}

export type AmazonModels = 'us.anthropic.claude-3-7-sonnet-20250219-v1:0'|'us.anthropic.claude-sonnet-4-20250514-v1:0';
export function narrowAmazonModel(model: any) : model is AmazonModels {
  return [
    'us.anthropic.claude-3-7-sonnet-20250219-v1:0',
    'us.anthropic.claude-sonnet-4-20250514-v1:0'
  ].includes(model);
}

export const MAX_TOKENS = {
  'claude-sonnet-4-20250514': 64000,
  'claude-3-7-sonnet-20250219': 64000,
  'claude-3-5-sonnet-latest': 8192,
  'claude-3-5-haiku-latest': 8192,
  'us.anthropic.claude-3-7-sonnet-20250219-v1:0': 64000,
  'us.anthropic.claude-sonnet-4-20250514-v1:0': 30000
}

// Minimal redefinition of the types from pluginAgent.ts that are needed here.
// (Re-export them so that derived classes can share them.)
export type AcceptedImageFormat = "image/jpeg" | "image/png";
export type UserPrompt =
  | { type: "text"; text: string }
  | {
      type: "image";
      source: {
        type: "base64";
        media_type: AcceptedImageFormat;
        data: string;
      };
    }
  | {
      type: "tool_result";
      tool_use_id: string;
      content: string;
    };;

function narrowFormat(format: string): format is AcceptedImageFormat {
  return ["image/jpeg", "image/png"].includes(format);
}

/**
 * A base Agent class for user prompt handling and chat orchestration.
 * Derived classes may override methods as needed.
 */
export abstract class BaseAgent {
  private mcpServerPath: string = "";
  private tools: string[] = [];
  private mcpClient: MCPClient|null = null;
  protected mcp: MCPClient | null = null;
  protected claude: Anthropic | null = null;
  protected openai: OpenAI | null = null;
  protected google: OpenAI | null = null;
  protected amazon: BedrockRuntimeClient | null = null;

  protected initialize() {
    if (!this.claude) {
      const rec = getAPIKey("claude");
      if (rec) {
        this.claude = new Anthropic({ apiKey: rec.apikey });
      }
    }
    if (!this.openai) {
      const rec = getAPIKey("openai");
      if (rec) {
        this.openai = new OpenAI({
          apiKey: rec.apikey, 
        });
      }
    }
    if (!this.google) {
      const rec = getAPIKey("google");
      if (rec) {
        this.google = new OpenAI({
          apiKey: rec.apikey,
          baseURL: "https://generativelanguage.googleapis.com/v1beta/openai/"
        });
      }
    }
    if (!this.amazon) {
      const aws_access_key_id = getAPIKey("aws_access_key_id")?.apikey; // process.env.AWS_ACCESS_KEY_ID;
      const aws_secret_access_key = getAPIKey("aws_secret_access_key")?.apikey; // process.env.AWS_SECRET_ACCESS_KEY;
      if (aws_access_key_id && aws_secret_access_key) {
        this.amazon = new BedrockRuntimeClient({
          region: "us-east-1",
          credentials: {
            accessKeyId: aws_access_key_id,
            secretAccessKey: aws_secret_access_key,
          },
        }); 
      }
    }
  }

  protected async addToolResultToChat(
    messages: ApptileInMemoryChatMessage[],
    toolInput: {
      partialInput?: string;
      tool: string;
      id: string;
    },
    textOfResult: string,
    temp_id: string,
    chatId?: number
  ) {
    messages.push({
      role: "user",
      content_type: "tool_result",
      tool_call_id: toolInput.id,
      tool_name: toolInput.tool,
      content: textOfResult,
      temp_id: temp_id
    });

    if (chatId) {
      await persistChatMessage(
        chatId,
        toolInput.id,
        null,
        "tool_result",
        "user",
        textOfResult,
        temp_id
      );
    } else {
      logger.error("No chatId found so tool response is not persisted!");
    }
  }

  protected async handleToolRun(
    toolInput: {
      partialInput?: string;
      tool: string;
      id: string;
    },
    messages: ApptileInMemoryChatMessage[],
    apptile_context: {
      app_id: string;
      plugin_root: string;
      plugins_dir: string;
      supabase_creds: null|SupabaseCreds;
    },
    temp_id: string,
    chatId?: number,
  ): Promise<NextStep> {
    let inputVars: any = {};
    if (toolInput.partialInput) {
      try {
        inputVars = JSON.parse(toolInput.partialInput);
      } catch (err) {
        console.error("Failed to parse toolinput: ", err);
      }
      inputVars.apptile_context = apptile_context;
    }

    if (!toolInput.id) {
      // force a tooluse id for gemini
      toolInput.id = "toolu-" + Date.now();
    }

    messages.push({
      role: "assistant",
      content_type: "tool_use",
      tool_call_id: toolInput.id,
      tool_name: toolInput.tool,
      content: toolInput.partialInput || "{}",
      temp_id: temp_id
    });

    if (chatId) {
      await persistChatMessage(
        chatId,
        toolInput.id,
        toolInput.tool,
        "tool_use",
        "assistant",
        toolInput.partialInput || "{}",
        temp_id
      );
    }

    if (this.mcp) {
      try {
        let result = await this.mcp.client.callTool({
          name: toolInput.tool,
          arguments: inputVars,
        });

        if (Array.isArray(result.content)) {
          if (result.content[0].text !== "user_response_required") {
            const textOfResult = result.content[0].text;

            await this.addToolResultToChat(messages, toolInput, textOfResult, temp_id, chatId);
            return "ContinueChat";
          } else {
            logger.info("Tool requires user intervention: " + toolInput?.tool);
            return "WaitForUserInput";
          } 
        } else {
          logger.error("Invalid response from tool. We are only working with one kind of resposne.");
          await this.addToolResultToChat(messages, toolInput, JSON.stringify(result), temp_id, chatId);
          return "WaitForUserInput";
        }
      } catch (err) {
        logger.error("Handled tool error happened", err);
        const error = err?.toString?.() || "There was an unrecognized error during tool execution";
        await this.addToolResultToChat(messages, toolInput, error, temp_id, chatId);
        return "ToolError";
      }
    } else {
      logger.error("No mcp client found to run the tool! " + toolInput.tool);
      await this.addToolResultToChat(messages, toolInput, "This tool is not available!", temp_id, chatId);
      return "ToolError";
    }
  }

  /**
   * Creates a minimal text/image prompt from user text input.
   * Derived classes may add custom logic or context if needed.
   */
  public prepareUserPrompt(
    userInput: string | {
      isToolResponse: boolean; 
      response: Array<{result: string; id: string;}>; 
    }, 
    assetName: string
  ): UserPrompt[] {
    let contextInfo = ""; 
    if (assetName) {
      contextInfo = `The name of the plugin you are working on is: '${assetName}'`;
    }

    let result: UserPrompt[] = [];
    if (typeof userInput === "string") {
      let imgTagStart = userInput.indexOf("<img");
      if (imgTagStart < 0) {
        imgTagStart = userInput.indexOf("&lt;img");
      }
      if (imgTagStart < 0) {
        result.push({ type: "text", text: userInput });
        if (contextInfo) {
          result.unshift({ type: "text", text: contextInfo });
        }
        return result;
      }

      const format = userInput.match(/img\ssrc="data:(.+);base64,/);
      if (format && format[1]) {
        const originalMessage = userInput;
        const dataStart = userInput.indexOf(";base64,") + ";base64,".length;
        const dataEnd = userInput.indexOf('"', dataStart);

        let imageTagEnd = userInput.indexOf(">", imgTagStart);
        if (imageTagEnd < 0) {
          imageTagEnd = userInput.indexOf("&gt;", imgTagStart);
        }
        userInput =
          userInput.slice(0, imgTagStart) +
          "`image_included_in_next_message`" +
          userInput.slice(imageTagEnd + 1, dataEnd);

        if (narrowFormat(format[1])) {
          result.push({ type: "text", text: userInput });
          result.push({
            type: "image",
            source: {
              type: "base64",
              media_type: format[1],
              data: originalMessage.slice(dataStart, dataEnd),
            },
          });
        } else {
          result.push({ type: "text", text: userInput });
        }
      } else {
        result.push({ type: "text", text: userInput });
      }

      if (contextInfo) {
        result.unshift({ type: "text", text: contextInfo });
      }
    } else if (userInput.isToolResponse) {
      result = userInput.response.map(resp => {
        return {
          type: "tool_result",
          tool_use_id: resp.id,
          content: resp.result
        };
      });
    } else {
      logger.error("Unprocessable user prompt! Stopping. " + JSON.stringify(userInput));
      throw new Error("Invalid prompt!");
    }
    return result;
  }

  private countOpenAITokens(
    messages: ChatCompletionMessageParam[],
    model: TiktokenModel
  ) {
    const encoder = encoding_for_model(model);

    let tokenCount = 0;
    for (const msg of messages) {
      tokenCount += encoder.encode(msg.role || "").length;
      if (msg.role === "assistant") {
        if (typeof msg.content === "string") {
          tokenCount += encoder.encode(msg.content).length;
        } else if (Array.isArray(msg.content)) {
          for (let block of msg.content) {
            if (block.type === "text") {
              tokenCount += encoder.encode(block.text).length;
            } else {
              throw new Error("Cannot encode non text tokens yet");
            }
          }
        }

        if (msg.tool_calls) {
          for (let toolCall of msg.tool_calls) {
            tokenCount += encoder.encode(toolCall.id).length;
            tokenCount += encoder.encode(toolCall.type).length;
            tokenCount += encoder.encode(toolCall.function.name).length;
            tokenCount += encoder.encode(toolCall.function.arguments).length;
          }
        }
      } else if (msg.role === "tool") {
        tokenCount += encoder.encode(msg.tool_call_id).length;
        if (typeof msg.content === "string") {
          tokenCount += encoder.encode(msg.content).length;
        } else {
          for (let block of msg.content) {
            if (block.type === "text") {
              tokenCount += encoder.encode(block.text).length;
            } else {
              throw new Error("Cannot encode non text tokens yet");
            }
          }
        }
      } else if (msg.role === "user") {
        if (typeof msg.content === "string") {
          tokenCount += encoder.encode(msg.content).length;
        } else {
          for (let block of msg.content) {
            if (block.type === "text") {
              tokenCount += encoder.encode(block.text).length;
            } else {
              throw new Error("Cannot encode non text tokens yet");
            }
          }
        }
      }
    }
    return tokenCount;
  }

  /**
   * Fits chat history within a specified token budget by removing oldest messages if needed.
   */
  public async fitHistoryWithinTokenBudget(
    messages: ApptileInMemoryChatMessage[],
    initialPrompt: string,
    model: string,
    provider: LLMProviders,
    pluginInfo: {pluginsDir: string, pluginName: string} | null,
    user: string,
  ): Promise<ApptileInMemoryChatMessage[]> {
    logger.info(`[BaseAgent][${user}] Start fitting with ` + messages.length + " messages.")
    const MAX_TOKENS_IN_HISTORY = 40000;

    if (provider === "claude") {
      if (!this.claude) {
        return messages;
      }

      if (pluginInfo === null) {
        let totalLength = messages.length;
        if (totalLength > 50) {
          messages = messages.slice(totalLength - 50);
        }
      } else {
        let newMessages = [];
        let startIndex = 0;
        if (messages.length > 15) {
          startIndex = messages.length - 15
        }
        logger.info("Starting from index: " + startIndex);

        for (let i = startIndex; i < messages.length; ++i) {
          const message = messages[i];
          if (message.tool_name === "apply_file_ops") {
            message.content = '{"message": "The body of this tool call was removed to save tokens"}';
          }
          newMessages.push(message);
        }

        messages = sequenceToolCallAndResponses(newMessages);
        try {
          logger.info(chalk.green("Adding file contents for: " + pluginInfo.pluginsDir + pluginInfo.pluginName));
          const formattedFileContents = await readAllPluginFilesByPluginName(pluginInfo.pluginsDir, pluginInfo.pluginName);
          const existingCode = INSERTED_CODE_PREFIX + '\n' + formattedFileContents.content[0].text; 
          messages = messages.filter(it => !it.content.startsWith(INSERTED_CODE_PREFIX))
          messages.unshift({
            role: "user",
            content_type: "text",
            content: existingCode,
            tool_call_id: null,
            tool_name: null,
            temp_id: "none"
          });
        } catch (err) {
          logger.error(`[BaseAgent] could not read plugin code`);
        }
      }

      // let prevTokens = 0;
      // let candidateMessages = messages.slice();
      // do {
      //   try {
      //     const tokenCount = await this.claude.messages.countTokens({
      //       model,
      //       system: initialPrompt,
      //       messages: toAnthropicMessages(candidateMessages),
      //     });
      //     prevTokens = tokenCount.input_tokens;

      //     if (prevTokens > MAX_TOKENS_IN_HISTORY && candidateMessages.length > 0) {
      //       let startIndex = Math.floor(candidateMessages.length / 2);
      //       if (startIndex === 0 && candidateMessages.length > 0) {
      //         startIndex = 1;
      //       }

      //       const messagesToDrop = candidateMessages.length - startIndex;
      //       candidateMessages = candidateMessages.slice(startIndex);
      //       candidateMessages = sequenceToolCallAndResponses(candidateMessages);
      //       logger.info(`[BaseAgent][${user}] prevTokens: `+ prevTokens);
      //       // logger.info(`[BaseAgent][${user}] Dropped message to fit token budget:`+ JSON.stringify(droppedMessage, null, 2));
      //       logger.info(`[BaseAgent][${user}] Dropped ${messagesToDrop} messages to fit token budget`);
      //     }
      //   } catch (err) {
      //     const droppedMessage = candidateMessages.shift();
      //     candidateMessages = sequenceToolCallAndResponses(candidateMessages);
      //     logger.error("An error occurred during fitting. Will drop a message: ");
      //     logger.info(`[BaseAgent][${user}] prevTokens: `+ prevTokens);
      //     // logger.info(`[BaseAgent][${user}] Dropped message to fit token budget:`+ JSON.stringify(droppedMessage, null, 2));
      //     logger.info(`[BaseAgent][${user}] Dropped message to fit token budget`);
      //   }
      // } while (prevTokens > MAX_TOKENS_IN_HISTORY && candidateMessages.length > 0);
      // messages = candidateMessages;
    } else if (provider === "openai") {
      if (!this.openai) {
        return messages;
      }
      let prevTokens = 0;
      do {
        const openaiMsgs = toOpenaiMessages(messages, initialPrompt);
        // TODO(gaurav): Handle narrowing the model. Right now its only a few so its fine.
        const tokenCount = this.countOpenAITokens(openaiMsgs, model as TiktokenModel);
        prevTokens = tokenCount;

        if (prevTokens > MAX_TOKENS_IN_HISTORY && messages.length > 0) {
          const droppedMessage = messages.shift();
          logger.info(`[BaseAgent][${user}] Dropped message to fit token budget: `, droppedMessage);
        }
      } while(prevTokens > MAX_TOKENS_IN_HISTORY && messages.length > 0);
    } else if (provider === "amazon") {
      if (pluginInfo === null) {
        let totalLength = messages.length;
        if (totalLength > 50) {
          messages = messages.slice(totalLength - 50);
        }
      } else {
        let newMessages = [];
        let startIndex = 0;
        if (messages.length > 15) {
          startIndex = messages.length - 15
        }
        logger.info("Starting from index: " + startIndex);

        for (let i = startIndex; i < messages.length; ++i) {
          const message = messages[i];
          if (message.tool_name === "apply_file_ops") {
            message.content = '{"message": "The body of this tool call was removed to save tokens"}';
          }
          newMessages.push(message);
        }

        messages = sequenceToolCallAndResponses(newMessages);
        try {
          logger.info(chalk.green("Adding file contents for: " + pluginInfo.pluginsDir + pluginInfo.pluginName));
          const formattedFileContents = await readAllPluginFilesByPluginName(pluginInfo.pluginsDir, pluginInfo.pluginName);
          const existingCode = INSERTED_CODE_PREFIX + '\n' + formattedFileContents.content[0].text; 
          messages = messages.filter(it => !it.content.startsWith(INSERTED_CODE_PREFIX))
          messages.unshift({
            role: "user",
            content_type: "text",
            content: existingCode,
            tool_call_id: null,
            tool_name: null,
            temp_id: "none"
          });
        } catch (err) {
          logger.error(`[BaseAgent] could not read plugin code`);
        }
      }
    } else {
      messages = sequenceToolCallAndResponses(messages);;
      logger.info(`[BaseAgent] no token counting available for gemini`);
    }

    return messages;
  }

  async retrieveOrCreateChat(
    anchor: string, 
    model: string, 
    type: AgentProgramType, 
    provider: LLMProviders,
    finalizedPrompt: UserPrompt[],
    systemPromptPath: string,
    appId: string,
    user: string,
    temp_id: string,
  ) {
    let chat  = await getChat(anchor);
    if (!chat) {
      logger.info(`[BaseAgent][${user}] creating chat for anchor because it wasn't found` + anchor);
      chat = await createChat(anchor, type, provider, model, appId);
    }

    if (chat) {
      logger.info(`[BaseAgent][${user}] db chat object found. Finalized prompt that needs to be added to chat history: ` + 
        chalk.gray(JSON.stringify(finalizedPrompt, null, 2))
      );

      for (let prompt of finalizedPrompt) {
        // logger.debug(chalk.yellow("Persisting: " + JSON.stringify(prompt, null, 2)))
        if (prompt.type === "text") {
          // logger.debug(chalk.yellow("Persisting as text "))
          await persistChatMessage(
            chat.id, 
            null,
            null,
            "text",
            "user", 
            prompt.text,
            temp_id
          );
        } else if (prompt.type === "image") {
          // logger.debug(chalk.yellow("Persisting as image "))
          await persistChatMessage(
            chat.id, 
            null,
            null,
            prompt.source.media_type,
            "user", 
            prompt.source.data,
            temp_id
          );
        } else if (prompt.type === "tool_result") {
          // logger.debug(chalk.yellow("Persisting as toolresult "))
          await persistChatMessage(
            chat.id, 
            prompt.tool_use_id,
            null,
            "tool_result",
            "user", 
            prompt.content,
            temp_id
          );
        } else {
          logger.error("Cannot persist message")
        }
      }
    } else {
      throw new Error("Could not retrieve or create chat");
    }

    const systemPrompt = await readFile(systemPromptPath, {encoding: "utf8"});

    // build the message array
    let messages: ApptileInMemoryChatMessage[];
    // gather chat history (excluding tool calls from user+assistant)
    if (chat) {
      const rawMessages = await getChatMessages(chat.id, "ASC");
      logger.info(`[BaseAgent][${user}] Retrieved ` + rawMessages.length + " messages after persisting prompt");
      const chatMessages = sequenceToolCallAndResponses(rawMessages);
      logger.info(`[BaseAgent][${user}] After resequencing ` + chatMessages.length + " remained");
      messages = toInMemoryChatMessage(chatMessages);
    } else {
      messages = [];
    }
    return {chat, messages, systemPrompt};
  }

  private async safeListTools(retries = 3) {
    for (let attempt = 0; attempt < retries; attempt++) {
      try {
        if (this.mcp) {
          return this.mcp?.client.listTools();
        } else {
          logger.error("[BaseAgent] Cannot list mcp tools as no mcp was found!");
          return {tools: []};
        }
      } catch (e) {
        console.error(`Attempt ${attempt + 1} failed:`, e);
        await this.reconnectMCPServer();
        await new Promise(r => setTimeout(r, 4000 * (attempt + 1)));
      }
    }
    logger.error("[BaseAgent] Failed to list mcp tools");
    return {tools: []}; // All retries failed
  }

  // Note: I have added the filtering on tools as a hack while the batch tool is evaluated in amazon
  protected async ingestResponseAndStream(
    llmsdk: Anthropic|OpenAI|BedrockRuntimeClient,
    systemPrompt: string,
    model: AnthropicModels|OpenAIModels|GoogleModels|AmazonModels,
    messages: ApptileInMemoryChatMessage[],
    res: Response,
    tokenCounterInstance: TokenCounter,
    isPlanner: boolean
  ): Promise<{
    status?: string;
    toolInput?: Array<{ partialInput?: string; tool: string; id: string }>;
    messageText?: string;
    uuid: string;
  } | null> {
    const uuid = v4();
    if (llmsdk instanceof Anthropic && narrowAnthropicModel(model)) {
      let tools: Anthropic.Tool[] = [];
      if (this.mcp) {
        const serverTools = await this.safeListTools();
        for (let i = 0; i < serverTools.tools.length; ++i) {
          const tool = serverTools.tools[i];
          if (isPlanner) {
            if (![
              "nocodelayer_create_global_plugin",
              "nocodelayer_create_screen",
              "nocodelayer_create_plugin",
              "nocodelayer_add_plugin_to_screen",
              "nocodelayer_update_global_plugin_defaultvalue"
            ].includes(tool.name)) {
              tools.push({
                name: tool.name,
                description: tool.description,
                input_schema: tool.inputSchema,
              })
            } 
          } else {
            if (tool.name !== "nocodelayer_batch_tools") {
              tools.push({
                name: tool.name,
                description: tool.description,
                input_schema: tool.inputSchema,
              })
            }
          }
        }
      }

      tokenCounterInstance.setTools(tools);

      if (tools.length > 0) {
        let lastTool = tools[tools.length - 1];
        lastTool.cache_control = {type: "ephemeral"};
      }

      const anthropicMessages = toAnthropicMessages(messages);
      logger.info(chalk.green("Calling anthropic"));
      const responseStream = llmsdk.messages.stream({
        system: [
          {
            type: "text",
            text: systemPrompt,
            cache_control: {type: "ephemeral"}
          }
        ],
        model,
        messages: anthropicMessages,
        max_tokens: MAX_TOKENS[model],
        tool_choice: {
          type: "auto",
          disable_parallel_tool_use: true,
        },
        tools,
      });
      logger.info(chalk.green("Stream started for temp_id" + uuid));

      const result = await ingestClaudeResponseStream(responseStream, res, uuid);
      logger.info(chalk.green("Stream processed for message: " + result.uuid));
      return result;
    } else if (llmsdk instanceof OpenAI && narrowOpenAIModel(model)) {
      let tools: ChatCompletionTool[] = [];
      if (this.mcp) {
        const serverTools = await this.safeListTools();
        tools = serverTools.tools.map(tool => {
          return {
            type: "function",
            function: {
              name: tool.name,
              description: tool.description,
              parameters: tool.inputSchema
            }
          } as any;
        })
        .filter(tool => tool.function.name !== "nocodelayer_batch_tools");
      }

      const openaiMessages = toOpenaiMessages(messages, systemPrompt);
      logger.info(chalk.green("Calling openai"));
      const responseStream = llmsdk.chat.completions.create({
        model,
        messages: openaiMessages,
        stream: true,
        tools,
        // temperature: 0.9
      });
      logger.info(chalk.green("Stream started"));

      const result = await ingestOpenaiMessageStream(responseStream, res, uuid);
      logger.info(chalk.green("Stream processed"));
      return result;
    } else if (llmsdk instanceof OpenAI && narrowGoogleModel(model)) {
      let tools: ChatCompletionTool[] = [];
      if (this.mcp) {
        const serverTools = await this.safeListTools();
        tools = serverTools.tools.map(tool => {
          return {
            type: "function",
            function: {
              name: tool.name,
              description: tool.description,
              parameters: tool.inputSchema,
            }
          } as any;
        })
        .filter(tool => tool.function.name !== "nocodelayer_batch_tools")
      }

      const openaiMessages = toOpenaiMessages(messages, systemPrompt);
      logger.info(chalk.green("Calling openai"));
      const responseStream = llmsdk.chat.completions.create({
        model,
        messages: openaiMessages,
        stream: true,
        tools,
        // temperature: 0.9
      });
      logger.info(chalk.green("Stream started"));

      const result = await ingestOpenaiMessageStream(responseStream, res, uuid);
      logger.info(chalk.green("Stream processed"));
      return result;
    } else if (llmsdk instanceof BedrockRuntimeClient && narrowAmazonModel(model)) {
      const tools: BedrockTool[] = [];
      if (this.mcp) {
        const serverTools = await this.safeListTools();
        for (let i = 0; i < serverTools.tools.length; ++i) {
          const tool = serverTools.tools[i];
          if (isPlanner) {
            if (![
              "nocodelayer_create_global_plugin",
              "nocodelayer_create_screen",
              "nocodelayer_create_plugin",
              "nocodelayer_add_plugin_to_screen",
              "nocodelayer_update_global_plugin_defaultvalue"
            ].includes(tool.name)) {
              tools.push({
                toolSpec: {
                  name: tool.name,
                  description: tool.description,
                  inputSchema: {
                    json: tool.inputSchema,
                  },
                }
              } as any)
            } 
          } else {
            if (tool.name !== "nocodelayer_batch_tools") {
              tools.push({
                toolSpec: {
                  name: tool.name,
                  description: tool.description,
                  inputSchema: {
                    json: tool.inputSchema,
                  },
                }
              } as any)
            }
          }
        }
      }

      const amazonMessages = toAmazonMessages(messages);
      logger.info(chalk.green("Calling amazon: ")); // + messages[0].content);

      const command = new ConverseStreamCommand({
        modelId: model,
        messages: amazonMessages,
        system: [{ text: systemPrompt }, { cachePoint: { type: "default"}  }],
        inferenceConfig: {
          maxTokens: 30000,
          temperature: 0.7,
          topP: 1,
        },
        toolConfig: {
          tools: tools
        }
      });

      logger.info(chalk.green("Stream started"));
      const responseStream = await llmsdk.send(command);
      // console.log(JSON.stringify(response, null, 4));
      
      const result = await ingestAmazonMessageStream(responseStream, res, uuid);
      logger.info(chalk.green("Stream processed"));
      if (result.amazonUsage) {
        tokenCounterInstance.setAmazonUsage(result.amazonUsage);
      }
      return result;
    } else {
      throw new Error("This llm provider is not implemented");
    }
  }

  /**
   * Default runPrompt method for Agents that rely on chat-based interactions.
   * Derived classes can override or extend it as needed.
   */
  public async runPrompt(
    appId: string,
    assetName: string,
    userPrompt: string,
    req: Request,
    res: Response,
    model: string,
    provider: LLMProviders,
    user: string,
    temp_id: string
  ): Promise<{
    tokenCounterInstance: Partial<TokenCounter> | null;
    message: string;
  }> {
    throw new Error("You must implement a runPrompt method in your class");
  }

  private async reconnectMCPServer() {
    if (this.mcpClient && this.mcpServerPath) {
      return this.connectToMCPServer(this.mcpServerPath, this.mcpClient, this.tools);
    } else {
      logger.error("Cannot reconnect mcp server before connection happens once!");
    }
  }

  /**
   * Helper method to connect to an MCP server and gather tools from it.
   */
  protected async connectToMCPServer(
    serverScriptPath: string,
    mcpClient: MCPClient,
    tools: string[] = []
  ) {
    this.mcpServerPath = serverScriptPath;
    this.mcpClient = mcpClient;
    this.tools = tools;

    logger.info(chalk.red("CONNECTING TO MCP SERVER: " + tools.join(",")))
    try {
      const isJS = serverScriptPath.endsWith(".js");
      const isPy = serverScriptPath.endsWith(".py");
      const isTS = serverScriptPath.endsWith(".ts");

      let transport;
      if (isJS) {
        transport = new StdioClientTransport({
          command: "node",
          args: [serverScriptPath, ...tools],
        });
      } else if (isPy) {
        transport = new StdioClientTransport({
          command: "python3",
          args: [serverScriptPath, ...tools],
        });
      } else if (isTS) {
        const tsNodePath = path.resolve(nodejsRoot, "node_modules", ".bin", "ts-node");
        transport =  new StdioClientTransport({
          command: "node",
          args: [tsNodePath, serverScriptPath, ...tools],
        });

      } else {
        transport = new StdioClientTransport({
          command: serverScriptPath,
          args: tools
        });
      }

      await mcpClient.client.connect(transport);

      // const toolResult = await mcpClient.client.listTools();
      // for (let i = 0; i < toolResult.tools.length; i++) {
      //   const tool = toolResult.tools[i];
      //   mcpClient.tools.push(tool.name);
      // }
    } catch (err) {
      logger.error("Failed to connect to mcp server: " + serverScriptPath, err);
    }
  }
}
