import { AIChat, ChatManager } from "./chat.js";
import { util } from "./util.js";
import { Session } from "./session.js";
import { log } from "./log.js";
import { ToolCall, ResourceCall } from "./type.js";
import { mcpModulesManager } from "./module.js";

class BaseAgent {
  session: Session;
  chatManager: ChatManager;

  /**
   * 取消当前聊天
   */
  cancelChat() {
    this.session.emit("cancel");
  }

  constructor() {
    this.session = new Session();

    this.chatManager = new ChatManager({
      session: this.session,
    });
  }

  async parseJson<T>(jsonString: string): Promise<T> {
    try {
      return JSON.parse(jsonString) as T;
    } catch (error: any) {
      const chatClient = await this.chatManager.createFixJsonChat();

      let content = await chatClient.chat(jsonString);
      // 处理大模型推理的think标签内容
      content = util.removeThinkConetnt(content);

      log.info("修复后的JSON内容", content);

      return util.parseJson(content) as T;
    }
  }
}

export class McpAgent extends BaseAgent {
  chatClient: AIChat | undefined;
  isChatting: boolean = false;

  constructor() {
    super();
  }

  async chat(message: string) {
    this.isChatting = true;
    try {
      await this.chatToCallFunction(message);
    } catch (error: any) {
      if (error.name === "AbortError") {
        return; // 聊天被取消，直接返回
      }
      // 只有非取消错误才记录日志并重新抛出
      throw error;
    } finally {
      this.isChatting = false;
    }
  }

  private async chatToCallFunction(usermessage: string) {
    let chatClient = this.chatClient;
    if (!chatClient) {
      chatClient = this.chatClient = await this.chatManager.createMcpChat();
    }

    // 多次重试原因说明：
    // 受限于大模型本身的能力，部分模型的推理结果可能存在错误，需要多次尝试
    // 如可能建议使用更好的模型，这样减少模型推理的错误，降低Token消耗
    const exec = async (usermessage: string, count: number) => {
      const countMax = 3; // 最大重试次数
      try {
        const content = await chatClient.chat(usermessage);
        // 处理工具调用
        const toolResult = await this.processToolOrResourceCall(content);

        if (toolResult === null) {
          this.session.emit("finish");
        } else {
          await this.chatToCallFunction(
            typeof toolResult === "string"
              ? toolResult
              : JSON.stringify(toolResult)
          );
        }

        return content;
      } catch (error: any) {
        if (error.name === "AbortError") {
          // 聊天被取消，直接抛出错误
          throw error;
        }
        if (count < countMax) {
          log.warn(`重试第 ${count + 1} 次: `, error);
          // 生成一段新的用户消息,将错误信息告知大模型，重新尝试
          usermessage = `请重新尝试，之前的错误信息是：${JSON.stringify(
            error
          )}`;
          return exec(usermessage, count + 1);
        } else {
          throw error;
        }
      }
    };
    await exec(usermessage, 0);

    this.session.add(chatClient.getMessages());
  }

  private async processToolOrResourceCall(content: string) {
    // 处理大模型推理的think标签内容
    content = util.removeThinkConetnt(content);

    // 判定是否是工具调用
    if (content.includes("<use_mcp_resource>")) {
      return this.handlerReadResource(content);
    } else if (content.includes("<use_mcp_tool>")) {
      return this.handlerCallTool(content);
    } else {
      return null;
    }
  }

  private async handlerReadResource(content: string) {
    // 解析 <use_mcp_resource> 标签的内容
    const resourceCallRegex =
      /<use_mcp_resource>([\s\S]*?)<\/use_mcp_resource>/;
    const exec = resourceCallRegex.exec(content);
    const resourceCallJsonString = exec ? exec[1].trim() : "";

    log.info("读取资源", resourceCallJsonString);

    if (!resourceCallJsonString) {
      return null;
    }
    const resourceCall: ResourceCall = await this.parseJson(
      resourceCallJsonString
    );

    const module = mcpModulesManager.getMcpModule({
      uuid: resourceCall.moduleId,
    });

    if (!module) {
      log.error(`Module with uuid ${resourceCall.moduleId} not found`);
      return null;
    }

    try {
      const result = await module.client.readResource({
        uri: resourceCall.uri,
      });

      log.info("资源读取返回结果", result);
      return result;
    } catch (error: any) {
      log.error(
        `Failed to read resource ${resourceCall.uri} from module ${resourceCall.moduleId}:`,
        error
      );
      return `Error reading resource: ${error.message || error}`;
    }
  }

  private async handlerCallTool(content: string) {
    const toolCallRegex = /<use_mcp_tool>([\s\S]*?)<\/use_mcp_tool>/;
    const exec = toolCallRegex.exec(content);
    const toolCallContent = exec ? exec[1].trim() : "";

    log.info("工具调用内容", toolCallContent);

    if (!toolCallContent) {
      return null;
    }
    const toolCall: ToolCall = await this.parseJson(toolCallContent);
    if (!toolCall || !toolCall.moduleId || !toolCall.tool) {
      log.error("Invalid tool call format.");
      return null;
    }

    const module = mcpModulesManager.getMcpModule({
      uuid: toolCall.moduleId,
    });

    if (!module) {
      log.error(`Module with uuid ${toolCall.moduleId} not found`);
      return null;
    }
    try {
      const toolResult = await module.client.callTool({
        name: toolCall.tool,
        arguments: toolCall.arguments || {},
      });

      log.info("工具调用返回结果", toolResult);
      return toolResult;
    } catch (error: any) {
      log.error(
        `Failed to call tool ${toolCall.tool} from module ${toolCall.moduleId}:`,
        error
      );
      return `Error calling tool: ${error?.message || error}`;
    }
  }
}