import { Chat, LMStudioClient, tool, type ChatMessage, type Tool } from "@lmstudio/sdk";
import { z } from "zod";
import { clamp, createWorkspace, stripReasoning, type Workspace } from "../workspace";
import { agentTools } from "./agent";
import { coreTools } from "./core";
import { editTools } from "./edit";
import { gitTools } from "./git";
import { projectTools } from "./project";
import { searchTools } from "./search";

declare const process: { env: Record<string, string | undefined> };

/**
 * A subagent is a second prediction against the same loaded model, with its own
 * short context and a deliberately narrowed set of tools. The point is context
 * isolation: the subagent burns twenty tool calls exploring and hands back five
 * lines, instead of filling the main conversation with directory listings.
 */
interface AgentType {
  key: string;
  summary: string;
  /** Capability overrides applied on top of the parent workspace. */
  powers: { write: boolean; shell: boolean; git: boolean };
  brief: string;
}

const AGENT_TYPES: AgentType[] = [
  {
    key: "explorer",
    summary: "Read-only. Finds things and explains how code works. Cannot change anything.",
    powers: { write: false, shell: false, git: false },
    brief:
      "You are an explorer subagent. You can read and search the workspace but you cannot change " +
      "anything, and you have no shell. Investigate thoroughly, then answer the question you were " +
      "given as compactly as you can: name the exact files and line numbers that matter and what " +
      "they do. Do not speculate about code you did not read.",
  },
  {
    key: "reviewer",
    summary: "Read-only plus git history. Reviews code and reports problems.",
    powers: { write: false, shell: true, git: false },
    brief:
      "You are a reviewer subagent. You can read, search and inspect git history, but you cannot " +
      "edit files. Review what you were asked to review and report concrete problems: file, line, " +
      "what is wrong, and why it matters. Rank them worst first. If you find nothing real, say so " +
      "plainly rather than inventing minor nitpicks.",
  },
  {
    key: "verifier",
    summary: "Runs the build and tests, reports what actually fails. Cannot edit files.",
    powers: { write: false, shell: true, git: false },
    brief:
      "You are a verifier subagent. Run the project's checks with the verify tool and report " +
      "exactly what passed and what failed, quoting the real error output. You cannot edit files, " +
      "so do not propose speculative rewrites -- your job is an accurate status report.",
  },
  {
    key: "worker",
    summary: "Full access. Carries out a small, well-defined change end to end.",
    powers: { write: true, shell: true, git: false },
    brief:
      "You are a worker subagent handling one specific task. Read before you edit, make the " +
      "smallest change that works, then call verify and fix what it reports. When you are done, " +
      "state which files you changed and what verify said. Stay strictly inside the task you were " +
      "given -- do not refactor anything else you notice.",
  },
];

/** Hard ceilings, since each round is a full prediction against a local model. */
const MAX_STEPS_CEILING = 25;
const DEFAULT_STEPS = 12;
let running = 0;
const MAX_CONCURRENT = 1;

function findType(key: string): AgentType | undefined {
  return AGENT_TYPES.find((entry) => entry.key === key.trim().toLowerCase());
}

/**
 * Builds the subagent's tool list from the same factories the main provider
 * uses, but against a workspace whose capabilities have been narrowed. A
 * read-only subagent therefore cannot be handed a write tool by accident -- the
 * factories simply return nothing. run_subagent itself is never included, so a
 * subagent cannot spawn its own.
 */
function toolsForType(parent: Workspace, type: AgentType): Tool[] {
  const scoped = createWorkspace({
    root: parent.root,
    maxFileSizeKb: parent.maxFileSizeKb,
    commandTimeoutSec: parent.commandTimeoutSec,
    allowWrite: parent.allowWrite && type.powers.write,
    allowShell: parent.allowShell && type.powers.shell,
    allowNetwork: false,
    allowGitWrite: parent.allowGitWrite && type.powers.git,
  });
  return [
    ...coreTools(scoped),
    ...projectTools(scoped),
    ...searchTools(scoped),
    ...editTools(scoped),
    ...agentTools(scoped),
    ...gitTools(scoped),
  ];
}

/** Exposed for tests: what a given subagent type would actually be handed. */
export function __toolsForType(parent: Workspace, key: string): Tool[] {
  const type = findType(key);
  return type === undefined ? [] : toolsForType(parent, type);
}

export interface SubagentOptions {
  /** The client LM Studio handed this plugin -- already authenticated. */
  client: LMStudioClient;
  /** Aborts the subagent when the user stops the parent generation. */
  abortSignal: AbortSignal;
  maxSteps: number;
  /** Model key to pin subagents to. Empty means "whatever is loaded". */
  modelKey: string;
}

/**
 * Picks the model a subagent runs on. `llm.model()` with no argument returns
 * *any* loaded model, which is wrong when several are loaded, so prefer an
 * explicit key and otherwise choose from the loaded list -- never triggering a
 * load of something the user did not ask for.
 */
async function resolveModel(options: SubagentOptions) {
  const key = options.modelKey.trim();
  if (key !== "") return options.client.llm.model(key);

  const loaded = await options.client.llm.listLoaded();
  if (loaded.length === 0) {
    throw new Error(
      "no model is loaded, so a subagent cannot run. Load a model, or do the work yourself.",
    );
  }
  return loaded[0];
}

async function describeModel(model: { getModelInfo: () => Promise<unknown> }): Promise<string> {
  try {
    const info = (await model.getModelInfo()) as { displayName?: string; modelKey?: string } | undefined;
    return info?.displayName ?? info?.modelKey ?? "unknown model";
  } catch {
    return "unknown model";
  }
}

export function subagentTools(ws: Workspace, options: SubagentOptions): Tool[] {
  const maxSteps = options.maxSteps;
  return [
    tool({
      name: "list_agent_types",
      description:
        "List the kinds of subagent you can delegate to with run_subagent, and what each is " +
        "allowed to do.",
      parameters: {},
      implementation: async () =>
        AGENT_TYPES.map((entry) => `${entry.key}: ${entry.summary}`).join("\n"),
    }),

    tool({
      name: "run_subagent",
      description:
        "Delegate a self-contained job to a separate agent that has its own fresh context and " +
        "reports back a short summary. Use this when a job would need many tool calls whose " +
        "details you do not need to keep -- searching a large codebase, reviewing a file, running " +
        "the test suite. You get the conclusion without the noise. Give the subagent the full " +
        "task in one message: it cannot see this conversation and cannot ask you questions.",
      parameters: {
        task: z
          .string()
          .describe(
            "The complete, self-contained instruction. State the goal, any file paths you already " +
              "know, and exactly what you want reported back.",
          ),
        agent_type: z
          .string()
          .default("explorer")
          .describe("One of: explorer, reviewer, verifier, worker. Call list_agent_types for details."),
        max_steps: z
          .number()
          .int()
          .min(1)
          .default(DEFAULT_STEPS)
          .describe("Maximum tool-calling rounds the subagent may take before it must answer."),
      },
      implementation: async ({ task, agent_type, max_steps }, ctx) => {
        const type = findType(agent_type);
        if (type === undefined) {
          return `Error: "${agent_type}" is not a subagent type. Available: ${AGENT_TYPES.map((t) => t.key).join(", ")}.`;
        }
        if (task.trim().length < 10) {
          return "Error: the task is too vague. The subagent cannot see this conversation, so spell out what it should do and what to report back.";
        }
        if (type.powers.write && !ws.allowWrite) {
          return `Error: a "${type.key}" subagent needs file-writing permission, which is switched off. Use explorer or reviewer, or ask the user to enable writing.`;
        }
        if (type.powers.shell && !ws.allowShell) {
          return `Error: a "${type.key}" subagent needs shell permission, which is switched off. Use an explorer subagent instead.`;
        }
        if (running >= MAX_CONCURRENT) {
          return "Error: a subagent is already running. Wait for it to finish before starting another.";
        }

        const steps = Math.min(max_steps, maxSteps, MAX_STEPS_CEILING);
        const tools = toolsForType(ws, type);
        const transcript: string[] = [];
        let lastAssistant = "";
        let toolCalls = 0;

        running++;
        const started = Date.now();
        try {
          const model = await resolveModel(options);
          const modelName = await describeModel(model);

          const chat = Chat.from([
            {
              role: "system",
              content:
                `${type.brief}\n\n` +
                `Workspace root: ${ws.root}. Only paths inside it exist.\n` +
                `You have at most ${steps} tool-calling rounds. Spend them; do not stop early, but ` +
                `do not waste them re-reading what you have already seen.\n` +
                `Your final message is the ONLY thing the agent that called you will see. Make it a ` +
                `complete answer that stands on its own, with no references back to this conversation.`,
            },
            { role: "user", content: task },
          ]);

          ctx.status(`Running ${type.key} subagent on ${modelName} (up to ${steps} rounds)`);
          await model.act(chat, tools, {
            maxPredictionRounds: steps,
            // Stopping the parent generation must stop the subagent too.
            signal: options.abortSignal,
            onMessage: (message: ChatMessage) => {
              const text = extractText(message);
              if (text === "") return;
              if (message.getRole?.() === "assistant") lastAssistant = text;
              transcript.push(text);
            },
            onToolCallRequestStart: () => {
              toolCalls++;
              ctx.status(`${type.key} subagent: ${toolCalls} tool call(s) so far`);
            },
          });

          const seconds = Math.round((Date.now() - started) / 1000);
          const report = lastAssistant.trim() === "" ? transcript.join("\n").trim() : lastAssistant.trim();
          if (report === "") {
            return `The ${type.key} subagent finished after ${toolCalls} tool call(s) without producing a report. Try a narrower task, or do it yourself.`;
          }
          return clamp(
            `--- report from ${type.key} subagent (${modelName}, ${toolCalls} tool calls, ${seconds}s) ---\n${report}`,
            8000,
            "subagent report",
          );
        } catch (error) {
          const message = (error as Error).message ?? String(error);
          if (options.abortSignal.aborted) {
            return "The subagent was cancelled before it finished.";
          }
          return `Error: the subagent failed -- ${message}. Do the work yourself with your own tools.`;
        } finally {
          running--;
        }
      },
    }),
  ];
}

/** ChatMessage shape varies by SDK version; get its text without assuming one. */
function extractText(message: unknown): string {
  if (typeof message === "string") return message;
  if (message === null || typeof message !== "object") return "";
  const candidate = message as { getText?: () => unknown; content?: unknown; text?: unknown };
  try {
    if (typeof candidate.getText === "function") {
      const value = candidate.getText();
      if (typeof value === "string") return stripReasoning(value);
    }
  } catch {
    // Fall through to the plain fields.
  }
  if (typeof candidate.text === "string") return stripReasoning(candidate.text);
  if (typeof candidate.content === "string") return stripReasoning(candidate.content);
  return "";
}

import { Chat, LMStudioClient, tool, type ChatMessage, type Tool } from "@lmstudio/sdk";
import { z } from "zod";
import { clamp, createWorkspace, stripReasoning, type Workspace } from "../workspace";
import { agentTools } from "./agent";
import { coreTools } from "./core";
import { editTools } from "./edit";
import { gitTools } from "./git";
import { projectTools } from "./project";
import { searchTools } from "./search";

declare const process: { env: Record<string, string | undefined> };

/**
 * A subagent is a second prediction against the same loaded model, with its own
 * short context and a deliberately narrowed set of tools. The point is context
 * isolation: the subagent burns twenty tool calls exploring and hands back five
 * lines, instead of filling the main conversation with directory listings.
 */
interface AgentType {
  key: string;
  summary: string;
  /** Capability overrides applied on top of the parent workspace. */
  powers: { write: boolean; shell: boolean; git: boolean };
  brief: string;
}

const AGENT_TYPES: AgentType[] = [
  {
    key: "explorer",
    summary: "Read-only. Finds things and explains how code works. Cannot change anything.",
    powers: { write: false, shell: false, git: false },
    brief:
      "You are an explorer subagent. You can read and search the workspace but you cannot change " +
      "anything, and you have no shell. Investigate thoroughly, then answer the question you were " +
      "given as compactly as you can: name the exact files and line numbers that matter and what " +
      "they do. Do not speculate about code you did not read.",
  },
  {
    key: "reviewer",
    summary: "Read-only plus git history. Reviews code and reports problems.",
    powers: { write: false, shell: true, git: false },
    brief:
      "You are a reviewer subagent. You can read, search and inspect git history, but you cannot " +
      "edit files. Review what you were asked to review and report concrete problems: file, line, " +
      "what is wrong, and why it matters. Rank them worst first. If you find nothing real, say so " +
      "plainly rather than inventing minor nitpicks.",
  },
  {
    key: "verifier",
    summary: "Runs the build and tests, reports what actually fails. Cannot edit files.",
    powers: { write: false, shell: true, git: false },
    brief:
      "You are a verifier subagent. Run the project's checks with the verify tool and report " +
      "exactly what passed and what failed, quoting the real error output. You cannot edit files, " +
      "so do not propose speculative rewrites -- your job is an accurate status report.",
  },
  {
    key: "worker",
    summary: "Full access. Carries out a small, well-defined change end to end.",
    powers: { write: true, shell: true, git: false },
    brief:
      "You are a worker subagent handling one specific task. Read before you edit, make the " +
      "smallest change that works, then call verify and fix what it reports. When you are done, " +
      "state which files you changed and what verify said. Stay strictly inside the task you were " +
      "given -- do not refactor anything else you notice.",
  },
];

/** Hard ceilings, since each round is a full prediction against a local model. */
const MAX_STEPS_CEILING = 25;
const DEFAULT_STEPS = 12;
let running = 0;
const MAX_CONCURRENT = 1;

function findType(key: string): AgentType | undefined {
  return AGENT_TYPES.find((entry) => entry.key === key.trim().toLowerCase());
}

/**
 * Builds the subagent's tool list from the same factories the main provider
 * uses, but against a workspace whose capabilities have been narrowed. A
 * read-only subagent therefore cannot be handed a write tool by accident -- the
 * factories simply return nothing. run_subagent itself is never included, so a
 * subagent cannot spawn its own.
 */
function toolsForType(parent: Workspace, type: AgentType): Tool[] {
  const scoped = createWorkspace({
    root: parent.root,
    maxFileSizeKb: parent.maxFileSizeKb,
    commandTimeoutSec: parent.commandTimeoutSec,
    allowWrite: parent.allowWrite && type.powers.write,
    allowShell: parent.allowShell && type.powers.shell,
    allowNetwork: false,
    allowGitWrite: parent.allowGitWrite && type.powers.git,
  });
  return [
    ...coreTools(scoped),
    ...projectTools(scoped),
    ...searchTools(scoped),
    ...editTools(scoped),
    ...agentTools(scoped),
    ...gitTools(scoped),
  ];
}

/** Exposed for tests: what a given subagent type would actually be handed. */
export function __toolsForType(parent: Workspace, key: string): Tool[] {
  const type = findType(key);
  return type === undefined ? [] : toolsForType(parent, type);
}

export interface SubagentOptions {
  /** The client LM Studio handed this plugin -- already authenticated. */
  client: LMStudioClient;
  /** Aborts the subagent when the user stops the parent generation. */
  abortSignal: AbortSignal;
  maxSteps: number;
  /** Model key to pin subagents to. Empty means "whatever is loaded". */
  modelKey: string;
}

/**
 * Picks the model a subagent runs on. `llm.model()` with no argument returns
 * *any* loaded model, which is wrong when several are loaded, so prefer an
 * explicit key and otherwise choose from the loaded list -- never triggering a
 * load of something the user did not ask for.
 */
async function resolveModel(options: SubagentOptions) {
  const key = options.modelKey.trim();
  if (key !== "") return options.client.llm.model(key);

  const loaded = await options.client.llm.listLoaded();
  if (loaded.length === 0) {
    throw new Error(
      "no model is loaded, so a subagent cannot run. Load a model, or do the work yourself.",
    );
  }
  return loaded[0];
}

async function describeModel(model: { getModelInfo: () => Promise<unknown> }): Promise<string> {
  try {
    const info = (await model.getModelInfo()) as { displayName?: string; modelKey?: string } | undefined;
    return info?.displayName ?? info?.modelKey ?? "unknown model";
  } catch {
    return "unknown model";
  }
}

export function subagentTools(ws: Workspace, options: SubagentOptions): Tool[] {
  const maxSteps = options.maxSteps;
  return [
    tool({
      name: "list_agent_types",
      description:
        "List the kinds of subagent you can delegate to with run_subagent, and what each is " +
        "allowed to do.",
      parameters: {},
      implementation: async () =>
        AGENT_TYPES.map((entry) => `${entry.key}: ${entry.summary}`).join("\n"),
    }),

    tool({
      name: "run_subagent",
      description:
        "Delegate a self-contained job to a separate agent that has its own fresh context and " +
        "reports back a short summary. Use this when a job would need many tool calls whose " +
        "details you do not need to keep -- searching a large codebase, reviewing a file, running " +
        "the test suite. You get the conclusion without the noise. Give the subagent the full " +
        "task in one message: it cannot see this conversation and cannot ask you questions.",
      parameters: {
        task: z
          .string()
          .describe(
            "The complete, self-contained instruction. State the goal, any file paths you already " +
              "know, and exactly what you want reported back.",
          ),
        agent_type: z
          .string()
          .default("explorer")
          .describe("One of: explorer, reviewer, verifier, worker. Call list_agent_types for details."),
        max_steps: z
          .number()
          .int()
          .min(1)
          .default(DEFAULT_STEPS)
          .describe("Maximum tool-calling rounds the subagent may take before it must answer."),
      },
      implementation: async ({ task, agent_type, max_steps }, ctx) => {
        const type = findType(agent_type);
        if (type === undefined) {
          return `Error: "${agent_type}" is not a subagent type. Available: ${AGENT_TYPES.map((t) => t.key).join(", ")}.`;
        }
        if (task.trim().length < 10) {
          return "Error: the task is too vague. The subagent cannot see this conversation, so spell out what it should do and what to report back.";
        }
        if (type.powers.write && !ws.allowWrite) {
          return `Error: a "${type.key}" subagent needs file-writing permission, which is switched off. Use explorer or reviewer, or ask the user to enable writing.`;
        }
        if (type.powers.shell && !ws.allowShell) {
          return `Error: a "${type.key}" subagent needs shell permission, which is switched off. Use an explorer subagent instead.`;
        }
        if (running >= MAX_CONCURRENT) {
          return "Error: a subagent is already running. Wait for it to finish before starting another.";
        }

        const steps = Math.min(max_steps, maxSteps, MAX_STEPS_CEILING);
        const tools = toolsForType(ws, type);
        const transcript: string[] = [];
        let lastAssistant = "";
        let toolCalls = 0;

        running++;
        const started = Date.now();
        try {
          const model = await resolveModel(options);
          const modelName = await describeModel(model);

          const chat = Chat.from([
            {
              role: "system",
              content:
                `${type.brief}\n\n` +
                `Workspace root: ${ws.root}. Only paths inside it exist.\n` +
                `You have at most ${steps} tool-calling rounds. Spend them; do not stop early, but ` +
                `do not waste them re-reading what you have already seen.\n` +
                `Your final message is the ONLY thing the agent that called you will see. Make it a ` +
                `complete answer that stands on its own, with no references back to this conversation.`,
            },
            { role: "user", content: task },
          ]);

          ctx.status(`Running ${type.key} subagent on ${modelName} (up to ${steps} rounds)`);
          await model.act(chat, tools, {
            maxPredictionRounds: steps,
            // Stopping the parent generation must stop the subagent too.
            signal: options.abortSignal,
            onMessage: (message: ChatMessage) => {
              const text = extractText(message);
              if (text === "") return;
              if (message.getRole?.() === "assistant") lastAssistant = text;
              transcript.push(text);
            },
            onToolCallRequestStart: () => {
              toolCalls++;
              ctx.status(`${type.key} subagent: ${toolCalls} tool call(s) so far`);
            },
          });

          const seconds = Math.round((Date.now() - started) / 1000);
          const report = lastAssistant.trim() === "" ? transcript.join("\n").trim() : lastAssistant.trim();
          if (report === "") {
            return `The ${type.key} subagent finished after ${toolCalls} tool call(s) without producing a report. Try a narrower task, or do it yourself.`;
          }
          return clamp(
            `--- report from ${type.key} subagent (${modelName}, ${toolCalls} tool calls, ${seconds}s) ---\n${report}`,
            8000,
            "subagent report",
          );
        } catch (error) {
          const message = (error as Error).message ?? String(error);
          if (options.abortSignal.aborted) {
            return "The subagent was cancelled before it finished.";
          }
          return `Error: the subagent failed -- ${message}. Do the work yourself with your own tools.`;
        } finally {
          running--;
        }
      },
    }),
  ];
}

/** ChatMessage shape varies by SDK version; get its text without assuming one. */
function extractText(message: unknown): string {
  if (typeof message === "string") return message;
  if (message === null || typeof message !== "object") return "";
  const candidate = message as { getText?: () => unknown; content?: unknown; text?: unknown };
  try {
    if (typeof candidate.getText === "function") {
      const value = candidate.getText();
      if (typeof value === "string") return stripReasoning(value);
    }
  } catch {
    // Fall through to the plain fields.
  }
  if (typeof candidate.text === "string") return stripReasoning(candidate.text);
  if (typeof candidate.content === "string") return stripReasoning(candidate.content);
  return "";
}