import path from "path";
import { Request, Response } from "express";
import { 
  getAPIKey, 
  getOpenApp, 
  AgentProgramType, 
  persistChatMessage,
  sequenceToolCallAndResponses, 
} from "../../database/cliconfig";

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

/**
 * We use a custom replacer for logging big strings in JSON
 */
export function logReplacer(key: string, value: any) {
  if (key === "data" && typeof value === "string" && value.length > 20) {
    return value.slice(0, 20) + "...[truncated]";
  }
  return value;
}

export const PlannerAgentToolGroups: ToolGroup[] = [
  "projectCodeReadForPlanner",
  "dashboardRead",
  // "dashboardWrite",
  "promptPluginAgent",
  "codeExecution",
  "supabaseTools",
  "supabaseReadTools",
  "batchedTool"
] as const;

export class AppPlannerAgent extends BaseAgent {
  constructor() {
    super();
    this.initialize().catch((err) => {
      console.error("Failed to initialize PluginAgent", err);
    });
  }

  async initialize() {
    super.initialize();

    if (!this.mcp) {
      this.mcp = {
        client: new Client({ name: "planner-agent-client", version: "1.0.0"}),
      };
      await this.connectToMCPServer(
        path.resolve(nodejsRoot, "controllers", "agents", "mcp-servers", "index.ts"),
        this.mcp,
        [
          "appPlannerAgent", // name of the mcp server. The rest are tools.
          ...PlannerAgentToolGroups
        ]
      );
    }
  }

  // public prepareUserPrompt(
  //   userInput: string | {
  //     isToolResponse: boolean; 
  //     response: Array<{result: string; id: string;}>; 
  //   }, 
  //   assetName: string): UserPrompt[] {
  //   if (typeof userInput === "string") {
  //     return super.prepareUserPrompt(userInput, assetName);
  //   } else {
  //     return userInput.response.map(resp => {
  //       return {
  //         type: "tool_result",
  //         tool_use_id: resp.id,
  //         content: resp.result
  //       };
  //     });
  //   }
  // }

  /**
   * 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(`[PlannerAgent][${user}][${appId}]`) + " runPrompt started with userPrompt: " + chalk.gray(JSON.stringify(userPrompt, null, 2)));

    let requestDropped = false;
    req.socket.on("close", () => {
      requestDropped = true;
      logger.info(chalkColor(`[PlannerAgent][${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(`[PlannerAgent][${user}]`) + " no anthropic client found!");
      return {
        tokenCounterInstance: null,
        message: "No anthropic client found!",
      };
    }

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

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

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

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

    const pluginsDir = path.resolve(
      openedApp.repoPath,
      "remoteCode",
      "plugins"
    );

    // we will build a file for writing
    logger.info(chalkColor(`[ArchAgent][${user}][${appId}]`) + " anchoring chat to " + appId);

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

    logger.info(chalkColor(`[ArchAgent][${user}][${appId}]`) + " Chat retrieved");

    // 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, 
      null,
      user
    );
    logger.info(chalkColor(`[ArchAgent][${user}][${appId}]`) + " messages fit inside token budget" + finalMessages.length);
    
    //message history array for counting tokens later
    const tokenCounterInstance = new TokenCounter([], "app", model, provider);
    try {
      let toolRun = false;
      do {
        logger.info(chalkColor(`[ArchAgent][${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,
          true
        );
        logger.info(
          chalkColor(`[ArchAgent][${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(`[ArchAgent][${user}][${appId}]`) + " Checking for token usage exceeeded error"); 

        if (ingestedResponse?.status === "MaxTokensUsed") {
          logger.info(
            chalkColor(`[ArchAgent][${user}][${appId}]`) + " Claude stopped in the middle. We'll add a line to continue and loop again."
          );
          await new Promise((resolve) => setTimeout(resolve, 20_000));
          finalMessages.push({
            role: "user",
            content_type: "text",
            tool_call_id: null,
            tool_name: null,
            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, 
            null, 
            user
          );
          logger.info(chalkColor(`[ArchAgent][${user}][${appId}]`) + " Continuing after refitting and pausing");
          continue; // back up and call again
        }

        // handle tool calls
        toolRun = !!ingestedResponse?.toolInput;
        if (toolRun && ingestedResponse?.toolInput) {
          //Putting the tool call in chat history message
          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: "none"
              },
              {
                role: "user",
                tool_call_id: toolcall.id,
                tool_name: toolcall.tool,
                content_type: "tool_result",
                content: "",
                temp_id: "none"
              },
            ])
          );
          logger.info(
            chalkColor(`[ArchAgent][${user}][${appId}]`) +
              " tool call detected. Processing."
          );
          
          let nextStep: NextStep = "ContinueChat";

          let supabaseCreds: SupabaseCreds | null = null;
          for (let toolcall of ingestedResponse.toolInput) {
            logger.info(
              chalkColor(`[ArchAgent][${user}]${appId}`) +
                " running tool: " +
                chalk.green(toolcall.tool)
            );
            if (toolcall.tool.startsWith("supabase_") && !supabaseCreds) {
              supabaseCreds = await supabaseConfigManager.getCreds(appId);
            }

            let res = await this.handleToolRun(
              toolcall,
              finalMessages,
              {
                plugin_root: "",
                plugins_dir: pluginsDir,
                app_id: appId,
                supabase_creds: supabaseCreds,
              },
              ingestedResponse.uuid,
              chat.id,
            );

            // once nextstep is set to waitforuserinput we don't unset it till user responds
            if (nextStep !== "WaitForUserInput") {
              nextStep = res;
            }
          }

          logger.info(
            chalkColor(`[ArchAgent][${user}][${appId}]`) +
              " Next step for tool: " +
              chalk.green(nextStep)
          );
          if (nextStep === "WaitForUserInput") {
            toolRun = false;
          } else if (nextStep === "ContinueChat" || nextStep === "ToolError") {
            toolRun = true;
          }
        } else {
          logger.info(
            chalkColor(`[ArchAgent][${user}][${appId}]`) +
              " Stopping because no toolrun was detected"
          );
        }
      } while (toolRun && !requestDropped);
    } catch (err: any) {
      console.error(chalk.red(`[ArchAgent][${user}][${appId}] Failure during Claude 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);
      }
      // res.write("\nAn error occurred\n");
      // res.write(err?.error?.message || "\nNo error details\n");
    } finally {
      // close the file
      // fileStream.end();
    }

    // res.write("turn finished\n");
    res.end();
    return {
      tokenCounterInstance,
      message: "Success",
    };
  }
}
