import path from "path";
import { Request, Response } from "express";
import { readdir } from "node:fs/promises";
import { sha256 } from "js-sha256";
import { 
  getOpenApp, 
  AgentProgramType, 
  persistChatMessage 
} from "../../database/cliconfig";
import { compileMultiplePlugins } from "../compiler";

import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import chalk from "chalk";
import { cycleColor, BaseAgent, AnthropicModels, AmazonModels, OpenAIModels, NextStep, LLMProviders, GoogleModels, INSERTED_CODE_PREFIX } from "./baseAgent";
import { nodejsRoot } from "../../projectrootpath";
import { logger } from "../../utils/logger";
import { ITokenCounterRequest, TokenCounter } from "../tokenCounter";
import { readAllPluginFilesByPluginName, ToolGroup } from "./mcp-servers/tools";
import { 
  supabaseConfigManager, 
  SupabaseCreds 
} from "./mcp-servers/supabase";
import OpenAI from "openai";
import Anthropic from "@anthropic-ai/sdk";
import { BedrockRuntimeClient } from "@aws-sdk/client-bedrock-runtime";
import * as Sentry from "@sentry/node";

export const PluginAgentToolGroups: ToolGroup[] = [
  "graphqlSchema",
  "projectCodeRead",
  "projectCodeWrite",
  "codeExecution",
  // "dashboardRead",
  "supabaseReadTools",
  "batchedPluginAgentTools"
] as const;
export class PluginAgent extends BaseAgent {
  constructor() {
    super();
    this.initialize().catch((err) => {
      logger.error("Failed to initialize PluginAgent", err);
    });
  }

  /**
   * Replaces the old 'initialize' logic. Sets up the claude client and the MCP clients.
   * The user can add more logic for supabase, with the commented code or additional servers.
   */
  async initialize() {
    super.initialize();

    // Set up GraphQL MCP if not connected yet:
    if (!this.mcp) {
      this.mcp = {
        client: new Client({ name: "plugin-agent-client", version: "1.0.0" }),
      };

      await this.connectToMCPServer(
        path.resolve(
          nodejsRoot,
          "controllers",
          "agents",
          "mcp-servers",
          "index.ts"
        ),
        this.mcp,
        [
          "pluginAgent", // name of the mcp server. 
          ...PluginAgentToolGroups
        ]
      );
    }
  }

  /**
   * After generation is complete, we finalize the code, parse out the code blocks,
   * then compile them with the old logic from "prompt()" function.
   */
  private async finalProcessingAndCompilation(openedAppRepoPath: string) {
    logger.info("compiling all plugins");
    await new Promise((resolve) => {
      setTimeout(() => {
        resolve({});
      }, 2000);
    });

    // Now compile
    const pluginsDir = path.resolve(openedAppRepoPath, "remoteCode/plugins");
    const entries = await readdir(pluginsDir, { withFileTypes: true });
    const plugins = entries
      .filter((it) => it.isDirectory())
      .map((it) => it.name);
    const compiledDirectoryName = sha256(plugins.join("_")).toString();
    return await compileMultiplePlugins(
      path.resolve(openedAppRepoPath, "remoteCode"),
      plugins,
      compiledDirectoryName
    );
  }

  /**
   * The main driver method that replaces the old 'prompt(...)' function. It orchestrates
   * the sub-steps using the class methods above.
   */
  public async runPrompt(
    appId: string,
    assetName: string,
    userPrompt: string | {
      isToolResponse: boolean; 
      response: Array<{result: string; id: string;}>; 
    },
    req: Request,
    res: Response,
    model: AnthropicModels | OpenAIModels | GoogleModels | AmazonModels,
    provider: LLMProviders,
    user: string,
    temp_id: string
  ): Promise<{
    tokenCounterInstance: Partial<TokenCounter> | null;
    message: string;
  }> {
    
    const chalkColor = cycleColor();
    logger.info(
      chalkColor(`[DevAgent][${user}][${appId}]`) + "runPrompt started with userPrompt: " +
      chalk.gray(JSON.stringify(userPrompt))
    );

    let requestDropped = false;
    req.socket.on("close", () => {
      requestDropped = true;
      logger.info(chalkColor(`[PluginAgent][${user}][${appId}] request was closed. Tool loop will be abandoned if its still running.`))
    });

    // Prepare prompt data
    const finalizedPrompt = this.prepareUserPrompt(userPrompt, assetName);

    // ensure anthropic is ready
    if (!this.claude) {
      logger.info(chalkColor(`[PluginAgent][${user}][${appId}]`) + " no anthropic client found!");
      return {
        tokenCounterInstance: null,
        message: "No anthropic client found!",
      };
    }

    if (!this.openai) {
      logger.info(`[PluginAgent][${user}][${appId}] no openai client found!`);
      return {
        tokenCounterInstance: null,
        message: "No openai client found!",
      };
    }

    if (!this.google) {
      logger.info(chalkColor(`[PluginAgent][${user}]`) + " no google client found!");
      return {
        tokenCounterInstance: null,
        message: "No google client found!",
      }
    }

    if (!this.amazon) {
      logger.info(chalkColor(`[PluginAgent][${user}]`) + " no google client found!");
      return {
        tokenCounterInstance: null,
        message: "No amazon client found!",
      }
    }

    // get the open app
    const openedApp = await getOpenApp(appId);
    if (!openedApp) {
      logger.info(chalkColor(`[PluginAgent][${user}][${appId}]`) + " no app found for appId:" + appId);
      return {
        tokenCounterInstance: null,
        message: "App not found!",
      };
    }

    // we will build a file for writing
    const chatAnchor = path.resolve(
      openedApp.repoPath,
      "remoteCode",
      "plugins",
      assetName,
      "source",
      "component.jsx"
    );
    logger.info(chalkColor(`[DevAgent][${user}[${appId}]]`) + " Getting chat for: " + chatAnchor);

    const { chat, messages, systemPrompt } = await this.retrieveOrCreateChat(
      chatAnchor,
      model,
      "widget",
      provider,
      finalizedPrompt,
      path.resolve(nodejsRoot, provider === "amazon" ? "systemPrompts/tokenoptimized/latest.html" : "systemPrompts/latest.html"),
      appId,
      user,
      temp_id
    );

    logger.info(chalkColor(`[DevAgent][${user}][${appId}]`) + " Chat retrieved. Truncating to fit within limit.");
    // const slicedMessages = messages.length <= 10 ? messages : messages.slice(messages.length - 10)
    const slicedMessages = messages;
    // Now ensure we are under token budget
    let finalMessages = await this.fitHistoryWithinTokenBudget(
      slicedMessages,
      systemPrompt,
      model,
      provider,
      {pluginName: assetName, pluginsDir: path.resolve(openedApp.repoPath, "remoteCode", "plugins")},
      user
    );
    logger.info(chalkColor(`[DevAgent][${user}][${appId}]`) + " Starting chat with " + finalMessages.length + " messages after fitting.");

    const tokenCounterInstance = new TokenCounter([], "widget", model, provider);
    try {
      let toolRun = false;
      do {
        logger.info(chalkColor(`[PluginAgent][${user}][${appId}]`) + " triggering completion");

        //Putting the input message in chat history message
        tokenCounterInstance.appendInMemoryInputMessages(finalMessages);
        let llmsdk : OpenAI|Anthropic|BedrockRuntimeClient;
        if (provider === "claude") {
          llmsdk = this.claude;
        } else if (provider === "google") {
          llmsdk = this.openai;
        } else if (provider === "amazon") {
          llmsdk = this.amazon;
        } else {
          llmsdk = this.openai;
        }
        const ingestedResponse = await this.ingestResponseAndStream(
          llmsdk,
          systemPrompt,
          model,
          finalMessages,
          res,
          tokenCounterInstance,
          false
        );

        logger.info(
          chalkColor(`[DevAgent][${user}][${appId}]`) + " completion finished"
        );

        if (
          ingestedResponse?.messageText &&
          ingestedResponse.messageText.trim()
        ) {
          //Pushing to chat history array
          tokenCounterInstance.appendInMemoryOutputMessages([
            {
              role: "assistant",
              tool_call_id: null,
              tool_name: null,
              content_type: "text",
              content: ingestedResponse?.messageText,
              temp_id: ingestedResponse.uuid
            },
          ]);

          await persistChatMessage(
            chat.id,
            null,
            null,
            "text",
            "assistant",
            ingestedResponse.messageText,
            ingestedResponse.uuid
          );
          finalMessages.push({
            role: "assistant",
            tool_call_id: null,
            tool_name: null,
            content_type: "text",
            content: ingestedResponse.messageText,
            temp_id: ingestedResponse.uuid
          });
        }
        logger.info(chalkColor(`[DevAgent][${user}][${appId}]`) + " Added model response to chat history");

        if (ingestedResponse?.status === "MaxTokensUsed") {
          logger.info(
            chalkColor(`[DevAgent][${user}][${appId}]`) + " MaxTokensUsed error encountered! Will truncate and wait 20sec before continuing."
          );
          res.write(
            "\nClaude used max tokens. We'll wait a sec and prompt to continue.\n"
          );
          await new Promise((resolve) => setTimeout(resolve, 20_000));
          finalMessages.push({
            role: "user",
            tool_call_id: null,
            tool_name: null,
            content_type: "text",
            content: `You stopped in the middle of generation. Please continue. If you were in the middle of calling the tool apply_file_ops
                then you can try splitting the ops across multiple tool calls to avoid running out of tokens.`,
            temp_id: "none"
          });
          finalMessages = await this.fitHistoryWithinTokenBudget(
            finalMessages,
            systemPrompt,
            model,
            provider,
            {pluginName: assetName, pluginsDir: path.resolve(openedApp.repoPath, "remoteCode", "plugins")},
            user
          );
          logger.info(chalkColor(`[DevAgent][${user}][${appId}]`) + " retrying after delay because max tokens were used")
          continue; // back up and call again
        }

        // handle tool calls
        toolRun = !!ingestedResponse?.toolInput;
        if (toolRun && ingestedResponse?.toolInput) {
          //Putting the tool call in chat history message
          //Putting the tool call output in chat history array
          tokenCounterInstance.appendInMemoryOutputMessages(
            ingestedResponse.toolInput?.flatMap((toolcall) => [
              {
                role: "assistant",
                tool_call_id: toolcall.id,
                tool_name: toolcall.tool,
                content_type: "tool_use",
                content: toolcall.partialInput || "",
                temp_id: ingestedResponse.uuid
              },
              {
                role: "user",
                tool_call_id: toolcall.id,
                tool_name: toolcall.tool,
                content_type: "tool_result",
                content: "",
                temp_id: ingestedResponse.uuid
              },
            ])
          );

          logger.info(chalkColor(`[PluginAgent][${user}]`) + "tool call detected");
          const pluginsDir = path.resolve(
            openedApp.repoPath,
            "remoteCode",
            "plugins"
          );
          const assetRoot = path.resolve(pluginsDir, assetName);
          try {
            let nextStep: NextStep = "ContinueChat";
            let supabaseCreds: SupabaseCreds | null = null;
            for (let toolcall of ingestedResponse.toolInput) {
              logger.info("Running tool: " + chalk.green(toolcall.tool));
              if (toolcall.tool.startsWith("supabase_") && !supabaseCreds) {
                supabaseCreds = await supabaseConfigManager.getCreds(appId);
              }
              let toolResult = await this.handleToolRun(
                toolcall,
                finalMessages,
                {
                  plugin_root: assetRoot,
                  plugins_dir: pluginsDir,
                  app_id: appId,
                  supabase_creds: supabaseCreds
                },
                ingestedResponse.uuid,
                chat.id
              );

              if (toolcall.tool === "compile_generated_code") {
                res.write("streamItemStart:" + JSON.stringify({type: "soft_refresh"}) + ":streamItemEnd")
              }

              // If even one tool requires user feedback then we stop
              if (nextStep !== "WaitForUserInput") {
                nextStep = toolResult;
              }
            }

            logger.info("Next step after tool run: " + chalk.green(nextStep));
            if (nextStep === "WaitForUserInput") {
              toolRun = false;
            } else if (nextStep === "ContinueChat") {
              toolRun = true;
            }
          } catch (err) {
            logger.error(chalk.red("Error in handleToolRun: "), err);
            finalMessages.push({
              role: "user",
              tool_call_id: null,
              tool_name: null,
              content_type: "text",
              content:
                `There was an error when performing the tool call for \
                        ${JSON.stringify(
                          ingestedResponse.toolInput
                        )}. If you were running \
                        file ops try doing one file op per tool call to prevent \
                        output token limit exceeded errors. Here is the error: ` +
                err?.toString(),
              temp_id: "none"
            });
            continue;
          }
        }
        // Update the current code
        if ((finalMessages.length > 0) && finalMessages[0].content.startsWith(INSERTED_CODE_PREFIX)) {
          try {
            const formattedFileContents = await readAllPluginFilesByPluginName(path.resolve(openedApp.repoPath, "remoteCode", "plugins"), assetName);
            const existingCode = INSERTED_CODE_PREFIX + '\n' + formattedFileContents.content[0].text;
            finalMessages[0].content = existingCode;
          } catch (err) {
            let errorMessage;
            if (err && err.toString) {
              errorMessage = err.toString();
              logger.error(chalk.red("Error in reading the contents of plugin code files") + errorMessage);
            } else {
              console.error(chalk.red("Error in reading the contents of plugin code files"), err);
            }
          }
        }

        if (!toolRun) {
          try {
            logger.info("Running compilation before exiting plugin agent!");
            const compilationResult = await this.finalProcessingAndCompilation(openedApp.repoPath)
            if (compilationResult.errors) {
              toolRun = true;
              finalMessages.push({
                role: "user",
                content_type: "text",
                content: "The generated code failed to compile. The errors are: " + JSON.stringify(compilationResult.errors),
                tool_call_id: null,
                tool_name: null,
                temp_id: null
              });
            } else {
              logger.info(compilationResult);
            }
          } catch (err) {
            logger.error("Failure during compilation: " + err);
          }
        }
      } while (toolRun && !requestDropped);
    } catch (err: any) {
      console.error(
        chalk.red(
          chalkColor(`[DevAgent][${user}][${appId}]`) + ` Failure during ${provider} conversation:`,
          err?.error?.error
        ),
        err
      );
      res.write(
        "streamItemStart:" +
          JSON.stringify({
            type: "error",
            details: err?.error?.error || "unexpected error " + err,
          }) +
          ":streamItemEnd"
      );

      if (err && process.env.ENABLE_SENTRY) {
        Sentry.captureException(err)
      }
    }
    // compile
    logger.info("Preparing to compile generated code");
    const compilationResult = await this.finalProcessingAndCompilation(
      openedApp.repoPath
    );
    if (compilationResult.errors) {
      res.write(
        "streamItemStart:" +
          JSON.stringify(
            {
              type: "compile_failed",
              details: compilationResult,
            },
            null,
            2
          ) +
          ":streamItemEnd"
      );
    } else {
      res.write(
        "streamItemStart:" +
          JSON.stringify(
            {
              type: "compile_successful",
              details: compilationResult,
            },
            null,
            2
          ) +
          ":streamItemEnd"
      );
    }
    res.end();
    return {
      tokenCounterInstance,
      message: "Success",
    };
  }
}
