import {
  createBashToolDefinition,
  createEditToolDefinition,
  createFindToolDefinition,
  createGrepToolDefinition,
  createLsToolDefinition,
  createReadToolDefinition,
  createWriteToolDefinition,
  defineTool,
  type ExtensionAPI,
  type ExtensionContext,
  type ToolDefinition,
  type ToolRenderResultOptions,
} from "@mariozechner/pi-coding-agent";
import { realpath } from "node:fs/promises";
import { resolve } from "node:path";
import { Type, type TSchema } from "typebox";
import { applyPatch } from "./src/apply-patch";
import { isPathInside } from "./src/path-utils";
import {
  detectModelProfile,
  prepareApplyPatchArgs,
  toEditArgs,
  toGrepArgs,
  toReadArgs,
  toShellCommand,
  toWriteArgs,
  type ModelProfile,
} from "./src/tool-args";

const CLAUDE_ALIAS_TOOLS = ["Read", "Edit", "Write", "Bash", "Grep", "Glob"];
const CODEX_ALIAS_TOOLS = ["shell_command", "apply_patch"];
const BUILTIN_TOOLS = ["read", "bash", "edit", "write", "grep", "find", "ls"] as const;
const LEGACY_MANAGED_TOOLS = ["LS"];
const MANAGED_TOOLS = new Set([...CLAUDE_ALIAS_TOOLS, ...CODEX_ALIAS_TOOLS, ...BUILTIN_TOOLS, ...LEGACY_MANAGED_TOOLS]);
const MAX_BUILTIN_DEFINITIONS = 64;

const CLAUDE_TOOL_DESCRIPTIONS = {
  Read: "Read a file from the local filesystem.",
  Edit: "A tool for editing files",
  Write: "Write a file to the local filesystem.",
  Bash: "Run shell command",
  Grep: `A powerful search tool built on ripgrep

  Usage:
  - ALWAYS use Grep for search tasks. NEVER invoke \`grep\` or \`rg\` as a Bash command. The Grep tool has been optimized for correct permissions and access.
  - Supports full regex syntax (e.g., "log.*Error", "function\\s+\\w+")
  - Filter files with glob parameter (e.g., "*.js", "**/*.tsx") or type parameter (e.g., "js", "py", "rust")
  - Output modes: "content" shows matching lines, "files_with_matches" shows only file paths (default), "count" shows match counts
  - Use Agent tool for open-ended searches requiring multiple rounds
  - Pattern syntax: Uses ripgrep (not grep) - literal braces need escaping (use \`interface\\{\\}\` to find \`interface{}\` in Go code)
  - Multiline matching: By default patterns match within single lines only. For cross-line patterns like \`struct \\{[\\s\\S]*?field\`, use \`multiline: true\`
`,
  Glob: `- Fast file pattern matching tool that works with any codebase size
- Supports glob patterns like "**/*.js" or "src/**/*.ts"
- Returns matching file paths sorted by modification time
- Use this tool when you need to find files by name patterns
- When you are doing an open ended search that may require multiple rounds of globbing and grepping, use the Agent tool instead`,
} as const;

const CODEX_SHELL_COMMAND_DESCRIPTION = `Runs a shell command and returns its output.
- Always set the \`workdir\` param when using the shell_command function. Do not use \`cd\` unless absolutely necessary.`;

const CODEX_APPLY_PATCH_DESCRIPTION = `## \`apply_patch\`

Use the \`apply_patch\` shell command to edit files.
Your patch language is a stripped‑down, file‑oriented diff format designed to be easy to parse and safe to apply. You can think of it as a high‑level envelope:

*** Begin Patch
[ one or more file sections ]
*** End Patch

Within that envelope, you get a sequence of file operations.
You MUST include a header to specify the action you are taking.
Each operation starts with one of three headers:

*** Add File: <path> - create a new file. Every following line is a + line (the initial contents).
*** Delete File: <path> - remove an existing file. Nothing follows.
*** Update File: <path> - patch an existing file in place (optionally with a rename).

May be immediately followed by *** Move to: <new path> if you want to rename the file.
Then one or more “hunks”, each introduced by @@ (optionally followed by a hunk header).
Within a hunk each line starts with:

For instructions on [context_before] and [context_after]:
- By default, show 3 lines of code immediately above and 3 lines immediately below each change. If a change is within 3 lines of a previous change, do NOT duplicate the first change’s [context_after] lines in the second change’s [context_before] lines.
- If 3 lines of context is insufficient to uniquely identify the snippet of code within the file, use the @@ operator to indicate the class or function to which the snippet belongs. For instance, we might have:
@@ class BaseClass
[3 lines of pre-context]
- [old_code]
+ [new_code]
[3 lines of post-context]

- If a code block is repeated so many times in a class or function such that even a single \`@@\` statement and 3 lines of context cannot uniquely identify the snippet of code, you can use multiple \`@@\` statements to jump to the right context. For instance:

@@ class BaseClass
@@ 	 def method():
[3 lines of pre-context]
- [old_code]
+ [new_code]
[3 lines of post-context]

The full grammar definition is below:
Patch := Begin { FileOp } End
Begin := "*** Begin Patch" NEWLINE
End := "*** End Patch" NEWLINE
FileOp := AddFile | DeleteFile | UpdateFile
AddFile := "*** Add File: " path NEWLINE { "+" line NEWLINE }
DeleteFile := "*** Delete File: " path NEWLINE
UpdateFile := "*** Update File: " path NEWLINE [ MoveTo ] { Hunk }
MoveTo := "*** Move to: " newPath NEWLINE
Hunk := "@@" [ header ] NEWLINE { HunkLine } [ "*** End of File" NEWLINE ]
HunkLine := (" " | "-" | "+") text NEWLINE

A full patch can combine several operations:

*** Begin Patch
*** Add File: hello.txt
+Hello world
*** Update File: src/app.py
*** Move to: src/main.py
@@ def greet():
-print("Hi")
+print("Hello, world!")
*** Delete File: obsolete.txt
*** End Patch

It is important to remember:

- You must include a header with your intended action (Add/Delete/Update)
- You must prefix new lines with \`+\` even when creating a new file
- File references can only be relative, NEVER ABSOLUTE.

You can invoke apply_patch like:

\`\`\`
shell {"command":["apply_patch","*** Begin Patch\\n*** Add File: hello.txt\\n+Hello, world!\\n*** End Patch\\n"]}
\`\`\`
`;

type ShellAliasInput = {
  command?: string;
  description?: string;
  cwd?: string;
  workdir?: string;
  timeout?: number;
  timeout_ms?: number;
  run_in_background?: boolean;
};

type BashToolDefinition = ReturnType<typeof createBashToolDefinition>;
type AnyToolDefinition = ToolDefinition<any, any, any>;
type BuiltinToolName = (typeof BUILTIN_TOOLS)[number];
type RenderContext = {
  args: any;
  cwd: string;
  [key: string]: any;
};
type RenderTheme = Parameters<NonNullable<BashToolDefinition["renderCall"]>>[1];

const builtinDefinitions = new Map<string, AnyToolDefinition>();

function timeoutSeconds(input: { timeout?: number; timeout_ms?: number }): number | undefined {
  if (typeof input.timeout === "number") return input.timeout;
  if (typeof input.timeout_ms === "number") return Math.max(1, Math.ceil(input.timeout_ms / 1000));
  return undefined;
}

function createBuiltinDefinition(name: BuiltinToolName, cwd: string): AnyToolDefinition {
  switch (name) {
    case "read":
      return createReadToolDefinition(cwd);
    case "edit":
      return createEditToolDefinition(cwd);
    case "write":
      return createWriteToolDefinition(cwd);
    case "bash":
      return createBashToolDefinition(cwd);
    case "grep":
      return createGrepToolDefinition(cwd);
    case "find":
      return createFindToolDefinition(cwd);
    case "ls":
      return createLsToolDefinition(cwd);
  }
}

function getRenderDefinition(name: BuiltinToolName, cwd: string): AnyToolDefinition {
  const key = `${name}:${cwd}`;
  const existing = builtinDefinitions.get(key);
  if (existing) {
    builtinDefinitions.delete(key);
    builtinDefinitions.set(key, existing);
    return existing;
  }
  const definition = createBuiltinDefinition(name, cwd);
  if (builtinDefinitions.size >= MAX_BUILTIN_DEFINITIONS) {
    const oldestKey = builtinDefinitions.keys().next().value;
    if (oldestKey) builtinDefinitions.delete(oldestKey);
  }
  builtinDefinitions.set(key, definition);
  return definition;
}

function toBashRenderArgs(input: ShellAliasInput): { command: string; timeout?: number } {
  return { command: toShellCommand(input), timeout: timeoutSeconds(input) };
}

function withRenderArgs(context: unknown, args: unknown): RenderContext {
  return { ...(context as RenderContext), args };
}

function replaceRenderedTitle(text: string, builtinName: BuiltinToolName, aliasLabel: string): string {
  return aliasLabel === builtinName ? text : text.replace(builtinName, aliasLabel);
}

function applyAliasToolTitle(component: unknown, builtinName: BuiltinToolName, aliasLabel: string): boolean {
  if (aliasLabel === builtinName || !component || typeof component !== "object") return false;
  const target = component as { text?: unknown; setText?: (text: string) => void; children?: unknown[]; invalidate?: () => void };
  if (typeof target.text === "string") {
    const next = replaceRenderedTitle(target.text, builtinName, aliasLabel);
    if (next !== target.text) {
      if (typeof target.setText === "function") target.setText(next);
      else target.text = next;
      return true;
    }
    return false;
  }
  let changed = false;
  for (const child of target.children ?? []) changed = applyAliasToolTitle(child, builtinName, aliasLabel) || changed;
  if (changed) target.invalidate?.();
  return changed;
}

function renderAliasCall(
  name: BuiltinToolName,
  args: unknown,
  theme: RenderTheme,
  context: unknown,
  aliasLabel: string = name,
) {
  const renderContext = context as RenderContext;
  const renderCall = getRenderDefinition(name, renderContext.cwd).renderCall;
  if (!renderCall) throw new Error(`${name} renderer is unavailable`);
  const component = renderCall(args, theme, withRenderArgs(renderContext, args) as any);
  applyAliasToolTitle(component, name, aliasLabel);
  return component;
}

function renderAliasResult(
  name: BuiltinToolName,
  args: unknown,
  result: unknown,
  options: ToolRenderResultOptions,
  theme: RenderTheme,
  context: unknown,
  aliasLabel: string = name,
) {
  const renderContext = context as RenderContext;
  const renderResult = getRenderDefinition(name, renderContext.cwd).renderResult;
  if (!renderResult) throw new Error(`${name} renderer is unavailable`);
  const component = renderResult(
    result as Parameters<NonNullable<AnyToolDefinition["renderResult"]>>[0],
    options,
    theme,
    withRenderArgs(renderContext, args) as any,
  );
  applyAliasToolTitle((renderContext.state as { callComponent?: unknown } | undefined)?.callComponent, name, aliasLabel);
  applyAliasToolTitle(component, name, aliasLabel);
  return component;
}

function renderShellAliasCall(args: ShellAliasInput, theme: RenderTheme, context: unknown) {
  return renderAliasCall("bash", toBashRenderArgs(args), theme, context);
}

function renderShellAliasResult(result: unknown, options: ToolRenderResultOptions, theme: RenderTheme, context: unknown) {
  const args = toBashRenderArgs((context as RenderContext).args as ShellAliasInput);
  return renderAliasResult("bash", args, result, options, theme, context);
}

async function resolveWorkdir(ctx: ExtensionContext, workdir: string | undefined): Promise<string | undefined> {
  if (!workdir) return undefined;
  const cwdRealPath = await realpath(ctx.cwd);
  const absolutePath = resolve(ctx.cwd, workdir);
  const workdirRealPath = await realpath(absolutePath);
  if (isPathInside(cwdRealPath, workdirRealPath)) return workdirRealPath;
  throw new Error(`Working directory escapes workspace: ${workdir}`);
}

type BuiltinAliasOptions<Input> = {
  name: string;
  label: string;
  description: string;
  parameters: TSchema;
  builtinName: BuiltinToolName;
  toArgs: (params: Input) => unknown;
  renderShell?: "self";
  validate?: (params: Input) => void;
};

function registerBuiltinAlias<Input>(pi: ExtensionAPI, options: BuiltinAliasOptions<Input>): void {
  pi.registerTool(
    defineTool({
      name: options.name,
      label: options.label,
      description: options.description,
      parameters: options.parameters,
      renderShell: options.renderShell,
      renderCall: (params, theme, context) =>
        renderAliasCall(options.builtinName, options.toArgs(params as Input), theme, context, options.label),
      renderResult: (result, renderOptions, theme, context) =>
        renderAliasResult(
          options.builtinName,
          options.toArgs((context as RenderContext).args as Input),
          result,
          renderOptions,
          theme,
          context,
          options.label,
        ),
      async execute(id, params, signal, onUpdate, ctx) {
        const input = params as Input;
        options.validate?.(input);
        return createBuiltinDefinition(options.builtinName, ctx.cwd).execute(
          id,
          options.toArgs(input),
          signal,
          onUpdate,
          ctx,
        );
      },
    }),
  );
}

function registerClaudeAliases(pi: ExtensionAPI): void {
  registerBuiltinAlias(pi, {
    name: "Read",
    label: "Read",
    description: CLAUDE_TOOL_DESCRIPTIONS.Read,
    builtinName: "read",
    parameters: Type.Object({
      file_path: Type.String({ description: "Path to the file to read" }),
      offset: Type.Optional(Type.Number({ description: "Line number to start reading from" })),
      limit: Type.Optional(Type.Number({ description: "Maximum number of lines to read" })),
    }),
    toArgs: toReadArgs,
  });

  registerBuiltinAlias(pi, {
    name: "Edit",
    label: "Edit",
    description: CLAUDE_TOOL_DESCRIPTIONS.Edit,
    builtinName: "edit",
    parameters: Type.Object({
      file_path: Type.String({ description: "Path to the file to edit" }),
      old_string: Type.String({ description: "Exact text to replace" }),
      new_string: Type.String({ description: "Replacement text" }),
      replace_all: Type.Optional(Type.Boolean({ description: "Not supported by this adapter" })),
    }),
    renderShell: "self",
    toArgs: toEditArgs,
    validate: (params: { replace_all?: boolean }) => {
      if (params.replace_all) throw new Error("Edit.replace_all is not supported by the Pi edit adapter");
    },
  });

  registerBuiltinAlias(pi, {
    name: "Write",
    label: "Write",
    description: CLAUDE_TOOL_DESCRIPTIONS.Write,
    builtinName: "write",
    parameters: Type.Object({
      file_path: Type.String({ description: "Path to the file to write" }),
      content: Type.String({ description: "Complete file contents" }),
    }),
    toArgs: toWriteArgs,
  });

  registerBuiltinAlias(pi, {
    name: "Bash",
    label: "Bash",
    description: CLAUDE_TOOL_DESCRIPTIONS.Bash,
    builtinName: "bash",
    parameters: Type.Object({
      command: Type.String({ description: "Command to execute" }),
      description: Type.Optional(Type.String({ description: "Short command description" })),
      timeout: Type.Optional(Type.Number({ description: "Timeout in seconds" })),
      run_in_background: Type.Optional(Type.Boolean({ description: "Runs through the shell when true" })),
    }),
    toArgs: toBashRenderArgs,
    validate: (params: { run_in_background?: boolean }) => {
      if (params.run_in_background) throw new Error("Bash.run_in_background is not supported by this adapter");
    },
  });

  registerBuiltinAlias(pi, {
    name: "Grep",
    label: "Grep",
    description: CLAUDE_TOOL_DESCRIPTIONS.Grep,
    builtinName: "grep",
    parameters: Type.Object({
      pattern: Type.String({ description: "Search pattern" }),
      path: Type.Optional(Type.String({ description: "File or directory to search" })),
      glob: Type.Optional(Type.String({ description: "Glob filter" })),
      case_sensitive: Type.Optional(Type.Boolean({ description: "Case-sensitive search" })),
      regex: Type.Optional(Type.Boolean({ description: "Treat pattern as regex" })),
      before_context: Type.Optional(Type.Number({ description: "Lines before each match" })),
      after_context: Type.Optional(Type.Number({ description: "Lines after each match" })),
      context: Type.Optional(Type.Number({ description: "Lines around each match" })),
      max_count: Type.Optional(Type.Number({ description: "Maximum matches" })),
      limit: Type.Optional(Type.Number({ description: "Maximum matches" })),
    }),
    toArgs: toGrepArgs,
  });

  registerBuiltinAlias(pi, {
    name: "Glob",
    label: "Glob",
    description: CLAUDE_TOOL_DESCRIPTIONS.Glob,
    builtinName: "find",
    parameters: Type.Object({
      pattern: Type.String({ description: "Glob pattern" }),
      path: Type.Optional(Type.String({ description: "Directory to search" })),
    }),
    toArgs: (params) => params,
  });

}

async function runShellAlias(
  id: string,
  input: ShellAliasInput,
  signal: AbortSignal | undefined,
  onUpdate: Parameters<ReturnType<typeof createBashToolDefinition>["execute"]>[3],
  ctx: ExtensionContext,
) {
  const command = toShellCommand(input);
  if (!command) throw new Error("Missing shell command");
  const workdir = await resolveWorkdir(ctx, input.cwd ?? input.workdir);
  const tool = createBashToolDefinition(
    ctx.cwd,
    workdir ? { spawnHook: (spawnContext) => ({ ...spawnContext, cwd: workdir }) } : undefined,
  );
  return tool.execute(id, { command, timeout: timeoutSeconds(input) }, signal, onUpdate, ctx);
}

function registerCodexAliases(pi: ExtensionAPI): void {
  const shellParameters = Type.Object({
    command: Type.String({ description: "The shell script to execute in the user's default shell" }),
    workdir: Type.Optional(Type.String({ description: "Working directory" })),
    cwd: Type.Optional(Type.String({ description: "Working directory" })),
    timeout_ms: Type.Optional(Type.Number({ description: "Timeout in milliseconds" })),
    timeout: Type.Optional(Type.Number({ description: "Timeout in seconds" })),
  });

  pi.registerTool(
    defineTool({
      name: "shell_command",
      label: "shell_command",
      description: CODEX_SHELL_COMMAND_DESCRIPTION,
      parameters: shellParameters,
      renderCall: renderShellAliasCall,
      renderResult: renderShellAliasResult,
      async execute(id, params, signal, onUpdate, ctx) {
        return runShellAlias(id, params, signal, onUpdate, ctx);
      },
    }),
  );

  pi.registerTool(
    defineTool({
      name: "apply_patch",
      label: "apply_patch",
      description: CODEX_APPLY_PATCH_DESCRIPTION,
      parameters: Type.Object({
        input: Type.String({ description: "Complete patch text, from Begin Patch through End Patch" }),
      }),
      prepareArguments: prepareApplyPatchArgs,
      executionMode: "sequential",
      async execute(_id, params, signal, _onUpdate, ctx) {
        const result = await applyPatch(ctx.cwd, params.input, signal);
        return { content: [{ type: "text", text: result.summary }], details: result };
      },
    }),
  );
}

function toolsForProfile(profile: ModelProfile): string[] | undefined {
  if (profile === "claude") return CLAUDE_ALIAS_TOOLS;
  if (profile === "codex") return CODEX_ALIAS_TOOLS;
  return undefined;
}

function sameTools(left: string[], right: string[]): boolean {
  return left.length === right.length && left.every((tool, index) => tool === right[index]);
}

function applyToolProfile(pi: ExtensionAPI, model: unknown, baseTools: string[]): void {
  const targetTools = toolsForProfile(detectModelProfile(model));
  const active = pi.getActiveTools();
  const preserved = active.filter((tool) => !MANAGED_TOOLS.has(tool));
  const next = targetTools ? [...preserved, ...targetTools] : [...preserved, ...baseTools];
  const nextTools = [...new Set(next)];
  if (!sameTools(active, nextTools)) pi.setActiveTools(nextTools);
}

export default function modelSuitableTools(pi: ExtensionAPI): void {
  registerClaudeAliases(pi);
  registerCodexAliases(pi);

  let baseTools: string[] = [];

  pi.on("session_start", (_event, ctx) => {
    baseTools = pi.getActiveTools().filter((tool) => !MANAGED_TOOLS.has(tool));
    applyToolProfile(pi, ctx.model, baseTools);
  });

  pi.on("model_select", (event) => {
    applyToolProfile(pi, event.model, baseTools);
  });
}

export { applyPatch, parseApplyPatch } from "./src/apply-patch";
export { detectModelProfile, prepareApplyPatchArgs, toGrepArgs, toShellCommand } from "./src/tool-args";