import {
  generateObject,
  NoSuchToolError,
  smoothStream,
  stepCountIs,
  streamText,
  type ModelMessage,
  type StepResult,
} from "ai";
import { chatPrompt } from "../prompts/chat";
import { providersRegistry } from "../providers";
import { logger } from "../utils/logger";

export interface ChatAgentCallbacks {
  onChunk?: (chunk: any) => void | Promise<void>;
  onError?: (error: any) => void | Promise<void>;
  onFinish?: (result: any) => void | Promise<void>;
  onStepFinish?: (step: StepResult<any>) => void | Promise<void>;
}

/**
 * Cleans messages by removing toolCall tags like [toolCall:call_id,tool-name]
 *
 * These tags are sometimes added by UI components or other parts of the system
 * to track tool calls, but they should not be sent to the AI model as they
 * can confuse the model or interfere with its responses.
 *
 * @param messages - Array of CoreMessage objects to clean
 * @returns Array of CoreMessage objects with toolCall tags removed from string content
 *
 * @example
 * ```ts
 * const messages = [
 *   { role: "user", content: "Hello [toolCall:call_123,edit-file] world" }
 * ];
 * const cleaned = cleanMessages(messages);
 * // Result: [{ role: "user", content: "Hello  world" }]
 * ```
 */
export const cleanMessages = (messages: ModelMessage[]): ModelMessage[] => {
  return messages.map((message) => {
    if (typeof message.content === "string") {
      // Remove toolCall tags using regex
      const cleanedContent = message.content.replace(
        /\[toolCall:[^\]]+\]/g,
        ""
      );
      return {
        ...message,
        content: cleanedContent,
      } as ModelMessage;
    }
    return message;
  });
};

const repairToolCall =
  (aiModel: any) =>
  async ({ toolCall, tools, parameterSchema, error }: any) => {
    if (NoSuchToolError.isInstance(error)) {
      return null; // do not attempt to fix invalid tool names
    }

    const tool = tools[toolCall.toolName as keyof typeof tools];

    const { object: repairedArgs } = await generateObject({
      model: aiModel,
      schema: tool.parameters,
      prompt: [
        `The model tried to call the tool "${toolCall.toolName}"` +
          ` with the following arguments:`,
        JSON.stringify(toolCall.args),
        `The tool accepts the following schema:`,
        JSON.stringify(parameterSchema(toolCall)),
        "Please fix the arguments.",
      ].join("\n"),
    });

    return { ...toolCall, args: JSON.stringify(repairedArgs) };
  };

/**
 * Main chat agent function that processes messages and returns a streaming response
 *
 * @param config - Configuration object containing model, messages, mode, and workingDir
 * @param options - Optional configuration like headers
 * @returns StreamText result for processing the chat
 */
export const chatAgent = async (
  {
    model,
    messages,
    mode,
    workingDir = process.cwd(),
    tools = {},
    callbacks,
  }: {
    model: string;
    messages: ModelMessage[];
    mode: "write" | "ask";
    workingDir: string;
    tools?: Record<string, any>;
    callbacks: ChatAgentCallbacks;
  },
  options: { headers?: Record<string, string>; abortSignal?: AbortSignal } = {}
) => {
  const [providerName = "openai", modelName = "gpt-4o-mini"] =
    model.split(":::");
  const aiModel = await providersRegistry.getModel(providerName, modelName);

  const total = Object.keys(tools).length;
  logger.debug(`🔧 Using ${total} tools provided by the controller.`);

  // Clean messages before sending to AI model to remove any toolCall tags
  // that might interfere with model responses or cause confusion
  const cleanedMessages = cleanMessages(messages);
  const {
    onFinish = null,
    onChunk = null,
    onError = null,
    onStepFinish = null,
  } = callbacks;
  return streamText({
    ...options,
    system: chatPrompt({
      hasTools: !!total,
      hasTool: (toolName: string) => Object.keys(tools).includes(toolName),
      mode,
      workingDir,
    }),
    tools: tools,
    model: aiModel,
    temperature: 0.3,
    abortSignal: options.abortSignal,
    experimental_repairToolCall: repairToolCall(aiModel),
    experimental_transform: smoothStream({
      delayInMs: 20, // optional: defaults to 10ms
      chunking: "line", // optional: defaults to 'word'
    }),
    // Stop after 5 steps
    stopWhen: stepCountIs(100),
    messages: cleanedMessages,
    onFinish: (result) => {
      if (onFinish) {
        onFinish(result);
      }
    },
    onChunk: (chunk) => {
      if (onChunk) {
        onChunk(chunk);
      }
    },
    onError: (error) => {
      if (onError) {
        onError(error);
      }
    },
    onStepFinish: (result) => {
      if (onStepFinish) {
        onStepFinish(result);
      }
    },
  });
};

/**
 * Backward-compatible wrapper for chatAgent for use in examples
 * This provides default values for the new required parameters
 */
export const chatAgentSimple = async (
  {
    model,
    messages,
    mode = "ask",
    workingDir = process.cwd(),
    tools = {},
  }: {
    model: string;
    messages: ModelMessage[];
    mode?: "write" | "ask";
    workingDir?: string;
    tools?: Record<string, any>;
  },
  options: { headers?: Record<string, string>; abortSignal?: AbortSignal } = {}
) => {
  return chatAgent(
    {
      model,
      messages,
      mode,
      workingDir,
      tools,
      callbacks: {},
    },
    options
  );
};