maestro

/**
 * @file secondaryAgent.ts
 * Secondary agent delegation tool: consult_secondary_agent.
 * Manual REST-level agent loop with JSON tool call parsing.
 */

import { tool, type Tool } from "@lmstudio/sdk";
import { writeFile, readFile, readdir, mkdir, rm } from "fs/promises";
import { join, dirname, isAbsolute } from "path";
import { z } from "zod";
import { validatePath, createSafeToolImplementation, type ToolContext } from "./shared";
import { getSharedInstances } from "../memory/toolsProvider";
import { resolveProjectMemoryTarget, buildProjectMemoryTags } from "../memory/projectMemory";

// --- Helpers ---

let _cachedSecondaryModel: string | null = null;
let _cachedEndpoint: string | null = null;

async function detectSecondaryModel(endpoint: string): Promise<string> {
  if (_cachedSecondaryModel && _cachedEndpoint === endpoint) return _cachedSecondaryModel;
  try {
    const res = await fetch(`${endpoint}/models`, { signal: AbortSignal.timeout(5_000) });
    if (!res.ok) return "local-model";
    const data = await res.json();
    const models: Array<{ id: string }> = data?.data ?? [];
    _cachedSecondaryModel = models.length >= 2 ? models[1].id : (models[0]?.id ?? "local-model");
    _cachedEndpoint = endpoint;
    return _cachedSecondaryModel;
  } catch { return "local-model"; }
}

/** Parse a tool call from model output (supports 3 JSON formats). */
function parseToolCall(content: string): { tool: string; args: Record<string, any> } | null {
  const trimmed = content.trim();
  const jsonMatch = trimmed.match(/\{[\s\S]*\}/);
  if (!jsonMatch) return null;
  try {
    const parsed = JSON.parse(jsonMatch[0]);
    if (parsed.tool && parsed.args) return parsed;
    if (parsed.name && parsed.arguments) {
      const args = parsed.arguments;
      if (parsed.name === "save_file") {
        if (args.path && !args.file_name) args.file_name = args.path;
        if (args.data && !args.content) args.content = args.data;
      }
      return { tool: parsed.name, args };
    }
    const toolNameMatch = trimmed.match(/to=([a-zA-Z0-9_.]+)/);
    if (toolNameMatch) {
      let toolName = toolNameMatch[1].replace(/^functions\./, "");
      let args = parsed;
      if (toolName === "save_file") {
        if (Array.isArray(args)) args = { files: args };
        if (args.path && !args.file_name) args.file_name = args.path;
        if (args.data && !args.content) args.content = args.data;
      }
      return { tool: toolName, args };
    }
  } catch { /* JSON parsing failed */ }
  return null;
}

const REFUSAL_KEYWORDS = [
  "i cannot browse", "i don't have access", "i can't access",
  "unable to browse", "real-time news", "no internet access",
  "as an ai", "i do not have the ability", "cannot access the internet",
];

function isRefusal(content: string): boolean {
  return REFUSAL_KEYWORDS.some(kw => content.toLowerCase().includes(kw));
}

// --- Sub-agent tool dispatch ---

interface DispatchContext {
  cwd: string;
  allowFileSystem: boolean;
  allowWeb: boolean;
  allowCode: boolean;
  originalRunJavascript: (p: { javascript: string }) => Promise<{ stdout: string; stderr: string }>;
  originalRunPython: (p: { python: string }) => Promise<{ stdout: string; stderr: string }>;
  pluginConfig: any;
}

async function dispatchTool(
  tc: { tool: string; args: Record<string, any> },
  dctx: DispatchContext,
  filesModified: string[],
): Promise<string> {
  // --- File System ---
  if (dctx.allowFileSystem) {
    if (tc.tool === "read_file" && tc.args?.file_name) {
      return await readFile(validatePath(dctx.cwd, tc.args.file_name), "utf-8");
    }
    if (tc.tool === "list_directory") {
      return JSON.stringify(await readdir(dctx.cwd));
    }
    if (tc.tool === "save_file") {
      if (Array.isArray(tc.args?.files)) {
        const saved: string[] = [];
        for (const f of tc.args.files) {
          const fName = f.file_name || f.name || f.path;
          const fContent = f.content || f.data;
          if (fName && fContent) {
            try {
              const fpath = validatePath(dctx.cwd, fName);
              await mkdir(dirname(fpath), { recursive: true });
              await writeFile(fpath, fContent, "utf-8");
              filesModified.push(fName); saved.push(fName);
            } catch { /* continue */ }
          }
        }
        return saved.length > 0 ? `Success: Saved ${saved.length} files: ${saved.join(", ")}` : "Error: No valid files in batch.";
      }
      const fileName = tc.args?.file_name || tc.args?.name || tc.args?.path;
      const content = tc.args?.content || tc.args?.data;
      if (fileName && content) {
        const fpath = validatePath(dctx.cwd, fileName);
        await mkdir(dirname(fpath), { recursive: true });
        await writeFile(fpath, content, "utf-8");
        filesModified.push(fileName);
        return `Success: File saved to ${fpath}`;
      }
      return "Error: Missing 'file_name' or 'content'.";
    }
    if (tc.tool === "replace_text_in_file" && tc.args?.file_name && tc.args?.old_string && tc.args?.new_string) {
      const fpath = validatePath(dctx.cwd, tc.args.file_name);
      const content = await readFile(fpath, "utf-8");
      if (!content.includes(tc.args.old_string)) return "Error: 'old_string' not found exactly.";
      const count = content.split(tc.args.old_string).length - 1;
      if (count > 1) return `Error: Found ${count} occurrences. Be more specific.`;
      await writeFile(fpath, content.replace(tc.args.old_string, tc.args.new_string), "utf-8");
      filesModified.push(tc.args.file_name);
      return "Success: Text replaced.";
    }
    if (tc.tool === "delete_files_by_pattern" && tc.args?.pattern) {
      if (tc.args.pattern.length > 100) throw new Error("Pattern too complex");
      const regex = new RegExp(tc.args.pattern);
      const start = Date.now();
      regex.test("safe_test_string_for_redos_check_1234567890_safe_test_string_for_redos_check_1234567890");
      if (Date.now() - start > 100) throw new Error("Pattern too slow");
      const files = await readdir(dctx.cwd);
      const deleted: string[] = [];
      for (const file of files) {
        if (regex.test(file)) { await rm(join(dctx.cwd, file), { force: true }); deleted.push(file); }
      }
      return `Deleted ${deleted.length} files: ${deleted.join(", ")}`;
    }
  }

  // --- Web ---
  if (dctx.allowWeb) {
    if (tc.tool === "wikipedia_search" && tc.args?.query) {
      try {
        const res = await fetch(`https://en.wikipedia.org/api/rest_v1/page/summary/${encodeURIComponent(tc.args.query)}`);
        if (res.ok) { const d = await res.json(); return JSON.stringify({ title: d.title, extract: d.extract }); }
        return "Wikipedia: no results found.";
      } catch { return "Wikipedia search failed."; }
    }
    if (tc.tool === "web_search" && tc.args?.query) {
      const { search, SafeSearchType } = await import("duck-duck-scrape");
      const r = await search(tc.args.query, { safeSearch: SafeSearchType.OFF });
      return JSON.stringify(r.results.slice(0, 3));
    }
    if (tc.tool === "fetch_web_content" && tc.args?.url) {
      let html = (await (await fetch(tc.args.url)).text());
      // Strip scripts, styles, nav, footer, and HTML tags — keep only readable text
      html = html
        .replace(/<script[\s\S]*?<\/script>/gi, "")
        .replace(/<style[\s\S]*?<\/style>/gi, "")
        .replace(/<nav[\s\S]*?<\/nav>/gi, "")
        .replace(/<footer[\s\S]*?<\/footer>/gi, "")
        .replace(/<header[\s\S]*?<\/header>/gi, "")
        .replace(/<[^>]+>/g, " ")
        .replace(/&[a-z]+;/gi, " ")
        .replace(/\s+/g, " ")
        .trim();
      return html.substring(0, 5000);
    }
  }

  // --- Code ---
  if (dctx.allowCode) {
    if (tc.tool === "run_python" && tc.args?.python) {
      const res = await dctx.originalRunPython({ python: tc.args.python });
      return res.stderr ? `Error: ${res.stderr}` : res.stdout;
    }
    if (tc.tool === "run_javascript" && tc.args?.javascript) {
      const res = await dctx.originalRunJavascript({ javascript: tc.args.javascript });
      return res.stderr ? `Error: ${res.stderr}` : res.stdout;
    }
  }

  // --- Memory ---
  if (tc.tool === "remember" && tc.args?.content) {
    try {
      const storagePath = dctx.pluginConfig.get("memoryStoragePath") || "";
      const memCfg = { activeProject: dctx.pluginConfig.get("activeProject") || "" };
      const { db: memDb, engine: memEngine } = await getSharedInstances(storagePath);
      const target = resolveProjectMemoryTarget(memCfg, { scope: tc.args.scope, project: tc.args.project });
      const tags = target.scope === "project" && target.project
        ? buildProjectMemoryTags(tc.args.tags || [], target.project)
        : (tc.args.tags || []);
      const id = memDb.store(tc.args.content, tc.args.category || "general", tags, 1.0, "sub-agent", null, target.scope, target.project);
      memEngine.indexMemory(id, tc.args.content, tags, tc.args.category || "general");
      return `Memory stored (id: ${id}, scope: ${target.scope}${target.project ? `, project: ${target.project}` : ""})`;
    } catch (err: any) { return `Memory store failed: ${err.message}`; }
  }
  if (tc.tool === "recall" && tc.args?.query) {
    try {
      const storagePath = dctx.pluginConfig.get("memoryStoragePath") || "";
      const { engine: memEngine } = await getSharedInstances(storagePath);
      const result = await memEngine.retrieve(tc.args.query, tc.args.limit || 5, 30);
      return JSON.stringify(result.memories.map(m => ({
        content: m.content, category: m.category, tags: m.tags,
        relevance: "compositeScore" in m ? Math.round((m as any).compositeScore * 100) : 50,
      })));
    } catch (err: any) { return `Memory recall failed: ${err.message}`; }
  }

  return "Error: Tool not found/allowed.";
}

// --- Auto-save code blocks from response ---

async function autoSaveCodeBlocks(
  finalContent: string,
  cwd: string,
  filesModified: string[],
): Promise<string> {
  const codeBlockRegex = /```\s*(\w+)?\s*([\s\S]*?)```/g;
  const matches = Array.from(finalContent.matchAll(codeBlockRegex));
  const processedFiles = new Set<string>();

  for (let i = matches.length - 1; i >= 0; i--) {
    const match = matches[i];
    const fullBlock = match[0], lang = (match[1] || "txt").toLowerCase(), code = match[2];
    const index = match.index || 0;
    let handledAsBatch = false;

    // Smart JSON Unpacking
    if (lang === "json") {
      try {
        const parsed = JSON.parse(code);
        if (Array.isArray(parsed)) {
          let extractedCount = 0;
          for (const item of parsed) {
            const fName = item.path || item.file_name || item.name;
            const fContent = item.content || item.data || item.code;
            if (fName && typeof fName === "string" && fContent && typeof fContent === "string") {
              const fpath = validatePath(cwd, fName);
              await mkdir(dirname(fpath), { recursive: true });
              await writeFile(fpath, fContent, "utf-8");
              filesModified.push(fName); processedFiles.add(fName); extractedCount++;
            }
          }
          if (extractedCount > 0) {
            handledAsBatch = true;
            finalContent = finalContent.slice(0, index) + `\n[System: Extracted ${extractedCount} files from JSON block.]\n` + finalContent.slice(index + fullBlock.length);
          }
        }
      } catch { /* not JSON batch */ }
    }

    if (!handledAsBatch && code.trim().length > 50) {
      const lookback = finalContent.substring(Math.max(0, index - 500), index);
      const EXT_PAT = /(?:tsx|ts|jsx|js|html|css|json|md|py|sh|java|rs|go|sql|yaml|yml|c|cpp|h|hpp|txt)/;
      const nameMatch = lookback.match(new RegExp(`(?:\`|\\*\\*|###|filename:|file:)[\\s\\S]*?([\\w\\-\\/\\\\.]+\\.(?:${EXT_PAT.source}))`, 'i'));
      let fileName = nameMatch?.[1]?.trim() || "";

      if (!fileName) {
        const firstLine = code.split('\n')[0].trim();
        const commentMatch = firstLine.match(new RegExp(`^(?:\\/\\/|#|<!--|;)\\s*(?:filename:|file:)?\\s*([\\w\\-\\/\\\\.]+\\.(?:${EXT_PAT.source}))`, 'i'));
        if (commentMatch) fileName = commentMatch[1].trim();
      }

      const isShell = ["bash", "sh", "cmd", "powershell", "console", "zsh", "terminal"].includes(lang);
      if ((isShell && !fileName) || !fileName || processedFiles.has(fileName)) continue;

      try {
        const fpath = validatePath(cwd, fileName);
        await mkdir(dirname(fpath), { recursive: true });
        await writeFile(fpath, code, "utf-8");
        filesModified.push(fileName); processedFiles.add(fileName);
        finalContent = finalContent.slice(0, index) + `\n[System: File '${fileName}' created successfully.]\n` + finalContent.slice(index + fullBlock.length);
      } catch { /* skip */ }
    }
  }
  return finalContent;
}

// --- Main export ---

export interface SecondaryAgentConfig {
  pluginConfig: any;
  originalRunJavascript: (p: { javascript: string }) => Promise<{ stdout: string; stderr: string }>;
  originalRunPython: (p: { python: string }) => Promise<{ stdout: string; stderr: string }>;
}

export function createSecondaryAgentTool(
  ctx: ToolContext,
  config: SecondaryAgentConfig,
): Tool {
  const { pluginConfig } = config;

  const _saFS = pluginConfig.get("subAgentAllowFileSystem");
  const _saWeb = pluginConfig.get("subAgentAllowWeb");
  const _saCode = pluginConfig.get("subAgentAllowCode");
  const _saCaps: string[] = ["memory (remember/recall)", "summarization"];
  if (_saFS) _saCaps.push("file reading");
  if (_saWeb) _saCaps.push("web search");
  if (_saCode) _saCaps.push("code execution");
  const _saNo: string[] = [];
  if (!_saCode) _saNo.push("coding", "file creation");
  if (!_saWeb) _saNo.push("web search");
  if (!_saFS) _saNo.push("file operations");
  const desc = `Delegate an auxiliary task to a secondary (lighter) model. Capabilities: ${_saCaps.join(", ")}.`
    + (_saNo.length > 0 ? ` Do NOT delegate ${_saNo.join(" or ")} — handle those yourself.` : "");

  return tool({
    name: "consult_secondary_agent",
    description: desc,
    parameters: {
      task: z.string().describe("The task to delegate."),
      agent_role: z.string().optional().describe("Key from 'Sub-Agent Profiles' config. Default: 'general'."),
      context: z.string().optional().describe("Additional context or data for the agent."),
      allow_tools: z.boolean().optional().describe("If true, the secondary agent can use its enabled tools. Default: false."),
    },
    implementation: createSafeToolImplementation(
      async ({ task, agent_role = "general", context = "", allow_tools = false }) => {
        const endpoint = pluginConfig.get("secondaryAgentEndpoint");
        const modelId = await detectSecondaryModel(endpoint);
        const subAgentProfilesStr = pluginConfig.get("subAgentProfiles");
        const debugMode = pluginConfig.get("enableDebugMode");
        const autoSave = pluginConfig.get("subAgentAutoSave");
        const showFullCode = pluginConfig.get("showFullCodeOutput");
        const allowFileSystem = pluginConfig.get("subAgentAllowFileSystem");
        const allowWeb = pluginConfig.get("subAgentAllowWeb");
        const allowCode = pluginConfig.get("subAgentAllowCode");

        const runAgentLoop = async (
          role: string, taskPrompt: string, contextData: string,
          loopLimit: number = 8, forceTools: boolean = false, workingDir: string,
        ) => {
          let systemPrompt = "You are a helpful assistant.";

          // Load instructions file if present
          try {
            const instructions = await readFile(join(workingDir, "SUB_AGENT_INSTRUCTIONS.md"), "utf-8");
            if (instructions.trim()) systemPrompt = instructions;
          } catch { /* ignore */ }

          systemPrompt += `\n\n## Current Workspace\nYour current working directory is:\n\n${workingDir}\nAlways assume relative paths are from this directory.`;

          // Append profile
          try {
            const profiles = JSON.parse(subAgentProfilesStr);
            if (profiles[role]) systemPrompt += `\n\n## Your Persona\n${profiles[role]}`;
            else if (role === "reviewer") systemPrompt += `\n\n## Your Persona\nYou are a Senior Code Reviewer. Analyze code, find bugs/issues, and FIX them using 'save_file'.`;
          } catch { /* ignore */ }

          // Append tools info
          let toolsReminder = "";
          const toolsEnabled = allow_tools || forceTools;
          if (toolsEnabled) {
            const allowedTools: string[] = [];
            if (allowFileSystem) allowedTools.push("read_file", "list_directory", "save_file", "replace_text_in_file", "delete_files_by_pattern");
            if (allowWeb) allowedTools.push("wikipedia_search", "web_search", "fetch_web_content");
            if (allowCode) allowedTools.push("run_python", "run_javascript");
            allowedTools.push("remember", "recall");
            if (allowedTools.length > 0) {
              systemPrompt += `\n\n## Allowed Tools\nYou have access to: ${allowedTools.join(", ")}.\n`;
              toolsReminder = `\n\n[SYSTEM REMINDER: You have access to tools: ${allowedTools.join(", ")}. USE A TOOL if needed.]`;
            }
          }

          const msgList = [
            { role: "system", content: systemPrompt },
            { role: "user", content: `Task: ${taskPrompt}\n\nContext: ${contextData}${toolsReminder}` },
          ];

          let loops = 0, finalContent = "";
          const filesModified: string[] = [];
          const dctx: DispatchContext = {
            cwd: workingDir, allowFileSystem, allowWeb, allowCode,
            originalRunJavascript: config.originalRunJavascript,
            originalRunPython: config.originalRunPython,
            pluginConfig,
          };

          while (loops < loopLimit) {
            try {
              const response = await fetch(`${endpoint}/chat/completions`, {
                method: "POST",
                headers: { "Content-Type": "application/json" },
                body: JSON.stringify({ model: modelId, messages: msgList, temperature: 0.7, stream: false }),
                signal: AbortSignal.timeout(120_000),
              });
              if (!response.ok) return { error: `API Error: ${response.status}`, filesModified };
              const data = await response.json();
              let content = data.choices[0].message.content;
              content = content.replace(/<\|.*?\|>/g, "").trim();

              if (!toolsEnabled) return { response: content, filesModified };

              if (isRefusal(content)) {
                msgList.push({ role: "assistant", content }, { role: "system", content: "SYSTEM ERROR: You HAVE access to tools. USE THEM." });
                loops++; continue;
              }

              const toolCall = parseToolCall(content);
              if (toolCall?.tool) {
                msgList.push({ role: "assistant", content });
                let toolResult: string;
                try { toolResult = await dispatchTool(toolCall, dctx, filesModified); }
                catch (err: any) { toolResult = `Error: ${err.message}`; }
                msgList.push({ role: "user", content: `Tool Output: ${toolResult}` });
                loops++;
              } else {
                if (content.includes("TASK_COMPLETED") || loops >= loopLimit - 1) {
                  finalContent = content; break;
                }
                msgList.push({ role: "assistant", content }, { role: "system", content: "SYSTEM NOTICE: You did not call a tool. If finished, output 'TASK_COMPLETED'. Otherwise USE A TOOL." });
                loops++;
              }
            } catch (err: any) { return { error: err.message, filesModified }; }

            // Prevent unbounded memory growth — keep system prompt + original task
            if (msgList.length > 20) {
              const sysMsg = msgList[0];
              const originalTask = msgList[1];
              msgList.splice(0, msgList.length, sysMsg, originalTask, ...msgList.slice(-17));
            }
          }

          // Auto-save code blocks
          if (autoSave && allowFileSystem && finalContent) {
            finalContent = await autoSaveCodeBlocks(finalContent, workingDir, filesModified);
          }

          return { response: finalContent, filesModified };
        };

        // --- 1. Primary Agent Loop ---
        const primaryResult = await runAgentLoop(agent_role, task, context, 8, false, ctx.cwd);
        if (primaryResult.error) return { error: primaryResult.error };
        let finalResponse = primaryResult.response;

        // --- 2. Auto-Debug Loop ---
        if (debugMode && allowCode && primaryResult.filesModified.length > 0) {
          const filesToCheck = primaryResult.filesModified.join(", ");
          let debugContext = "Here is the content of the created files:\n";
          for (const f of primaryResult.filesModified) {
            try { debugContext += `\n--- ${f} ---\n${await readFile(join(ctx.cwd, f), "utf-8")}\n`; } catch { /* skip */ }
          }
          if (debugContext.length > 8000) debugContext = debugContext.substring(0, 8000) + "\n\n[... truncated]";
          const debugResult = await runAgentLoop("reviewer", `Review the code in these files: ${filesToCheck}. Check for bugs, syntax errors, or logic flaws. Fix them.`, debugContext, 5, true, ctx.cwd);
          finalResponse += "\n\n--- Auto-Debug Report ---\n" + (debugResult.response || "Debug pass completed.");
          if (debugResult.filesModified.length > 0) finalResponse += `\n(Reviewer fixed: ${debugResult.filesModified.join(", ")})`;
        }

        // Append file list
        if (primaryResult.filesModified.length > 0) {
          const fullPaths = primaryResult.filesModified.map(f => isAbsolute(f) ? f : join(ctx.cwd, f));
          finalResponse += `\n\n[GENERATED_FILES]: ${fullPaths.join(", ")}`;
          if (showFullCode) {
            finalResponse += `\n\n### Generated Code Content:\n`;
            for (const f of primaryResult.filesModified) {
              try {
                const fpath = isAbsolute(f) ? f : join(ctx.cwd, f);
                const content = await readFile(fpath, "utf-8");
                finalResponse += `\n**${f}**\n\`\`\`${f.split('.').pop() || 'txt'}\n${content}\n\`\`\`\n`;
              } catch { /* skip */ }
            }
          }
        }

        if (!showFullCode) {
          finalResponse = finalResponse.replace(/```[\s\S]*?```/g, "\n[System: Code Block Hidden. The code has been handled by the sub-agent.]\n");
        }

        return { response: finalResponse, generated_files: primaryResult.filesModified };
      },
      pluginConfig.get("enableSecondaryAgent"),
      "consult_secondary_agent",
    ),
  });
}