// The model's tools. Offered only in a chat bound to a project; everywhere else the list is
// empty, so ordinary chats pay nothing. Each result is a few lines: what changed, where the
// plan stands, and what to do next. The full state is never repeated in a tool result.

import { tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { chatIdOf, loadChatState } from "./chat";
import { configSchematics } from "./config";
import { formatResults, hasIndex, search } from "./recall";
import { currentStep, Ledger, loadLedger, progress, readEvents, renderMarkdown, stepLabel } from "./ledger";
import { addDecision, addFinding, addQuestion, addSteps, nextLine, OpError, resolveQuestion, updateStep } from "./ops";
import { clip } from "./text";
import { absorbPendingEdit } from "./checks";

const STATUS = ["todo", "doing", "done", "blocked"] as const;
const VIEW_MAX_CHARS = 14000;

function nowLine(l: Ledger): string {
  const p = progress(l);
  const cur = currentStep(l);
  if (!p.total) return "now: no plan yet";
  return `now: ${cur ? clip(stepLabel(cur), 160) : "no step in progress"} | ${p.done} of ${p.total} done`;
}

function result(l: Ledger, changed: string, next?: string): string {
  return [`status: ok`, `changed: ${changed}`, nowLine(l), next ?? nextLine(l)].join("\n");
}

const RUN_AHEAD =
  "next: The plan was finished earlier in this reply, so these steps go beyond what was asked. Unless the user's latest message asked for them, give the user the result now, name these steps as what could come next, and stop.";

const NEXT_STEP_SAME_REPLY =
  "next: A step was finished earlier in this reply. If the user's latest message asked only for that step, give the user its result now and stop; start this step when they ask for it. An automatic note from context management telling you to continue is not a new request from the user.";

// What happened to the plan during this reply, from the event log: the steps the model added
// after finishing the whole plan, and whether it finished any step at all. Local models tend
// to keep going through a list the user means to take one message at a time, especially
// after context compression tells them to "continue the task"; the tool result is the one
// place to say so at the moment it happens.
export function replyProgress(l: Ledger, turnAt: string | undefined): { beyondPlan: Set<number>; finished: Set<number> } {
  const beyondPlan = new Set<number>();
  const finished = new Set<number>();
  if (!turnAt) return { beyondPlan, finished };
  let planDone = false;
  const plan = new Map<string, string>();
  for (const e of readEvents(l.name, 0)) {
    if (e.op === "step_added" && e.id) plan.set(e.id, "todo");
    if (e.op === "step_updated" && e.id) {
      const m = /now (todo|doing|done|blocked)/.exec(e.text ?? "");
      if (m) plan.set(e.id, m[1]);
    }
    if (e.ts < turnAt) continue;
    if (e.op === "step_updated" && e.id && /now done/.test(e.text ?? "")) {
      finished.add(Number(e.id));
      if ([...plan.values()].every((s) => s === "done")) planDone = true;
    } else if (planDone && e.op === "step_added" && e.id) beyondPlan.add(Number(e.id));
  }
  return { beyondPlan, finished };
}

function failure(message: string, next: string): string {
  return [`status: error`, `reason: ${message}`, `next: ${next}`].join("\n");
}

// Every tool reloads the ledger when called, so changes made by commands or other chats are seen.
function withLedger(slug: string, fn: (l: Ledger) => string): string {
  const l = loadLedger(slug);
  if (!l) return failure("This chat's project no longer exists.", "Tell the user; they can start or open a project with /tl.");
  try {
    // The user may have saved an edit to ledger.md since the last message: merge it first.
    absorbPendingEdit(l);
    return result(l, fn(l));
  } catch (e) {
    if (e instanceof OpError) return failure(e.message, "Fix the call and try again, or continue without it.");
    return failure("The ledger could not be written.", "Continue without recording, and tell the user if it keeps failing.");
  }
}

function viewSection(l: Ledger, section: string): string {
  const md = renderMarkdown(l);
  if (section === "all") return md;
  const heading: Record<string, string> = {
    goal: "## Goal", constraints: "## Constraints", plan: "## Plan", findings: "## Findings", decisions: "## Decisions", questions: "## Open questions",
  };
  const start = md.indexOf(heading[section]);
  if (start < 0) return md;
  const rest = md.slice(start);
  const end = rest.indexOf("\n## ", 3);
  return end > 0 ? rest.slice(0, end) : rest;
}

export function ledgerTools(chatId: string, slug: string): Tool[] {
  return [
    tool({
      name: "ledger_view",
      description:
        "Show this chat's project ledger: goal, constraints, plan, findings with sources, decisions and questions. Use it when the throughLine state in the user's message is not enough.",
      parameters: { section: z.enum(["all", "goal", "constraints", "plan", "findings", "decisions", "questions"]).optional() },
      implementation: async ({ section }) => {
        const l = loadLedger(slug);
        if (!l) return failure("This chat's project no longer exists.", "Tell the user.");
        let text = viewSection(l, section ?? "all");
        if (text.length > VIEW_MAX_CHARS) text = text.slice(0, VIEW_MAX_CHARS) + "\n(cut here; view one section at a time for the rest)";
        return `status: ok\n${text}\n${nextLine(l)}`;
      },
    }),
    tool({
      name: "ledger_plan",
      description:
        "Add steps to the project's plan, or change a step's status, note or wording. Plan only what the goal needs; add steps beyond it only if the user asks. Keep exactly one step 'doing' while you work, and mark it 'done' when finished.",
      parameters: {
        action: z.enum(["add", "update"]),
        steps: z.array(z.string()).optional().describe("For add: one or more step texts, in order."),
        step: z.number().int().optional().describe("For update: the step number."),
        status: z.enum(STATUS).optional(),
        note: z.string().optional().describe("A short note on the step, such as what is left or why it is blocked."),
        text: z.string().optional().describe("For update: new wording for the step."),
      },
      implementation: async ({ action, steps, step, status, note, text }) => {
        let warning: string | null = null;
        const out = withLedger(slug, (l) => {
          const turnAt = loadChatState(chatId)?.turnAt;
          if (action === "add") {
            const changed = addSteps(l, steps?.length ? steps : text ? [text] : [], status ?? "todo", "model", chatId);
            if (replyProgress(l, turnAt).beyondPlan.has(l.plan.at(-1)!.id)) warning = RUN_AHEAD;
            return changed;
          }
          if (step === undefined) throw new OpError("For update, give the step number.");
          const changed = updateStep(l, step, { status, note, text }, "model", chatId);
          if (status === "doing") {
            const p = replyProgress(l, turnAt);
            if (p.beyondPlan.has(step)) warning = RUN_AHEAD;
            else if ([...p.finished].some((id) => id !== step)) warning = NEXT_STEP_SAME_REPLY;
          }
          return changed;
        });
        return warning ? out.replace(/\nnext: .*$/, "\n" + warning) : out;
      },
    }),
    tool({
      name: "ledger_record",
      description:
        "Record something established for the project: a finding (a fact, with the source it came from), a decision (with the reason), or an open question that blocks progress. Project facts only; facts about the user belong in memory.",
      parameters: {
        kind: z.enum(["finding", "decision", "question"]),
        text: z.string(),
        source: z.string().optional().describe("For a finding: the URL of the page you read it on, a file path, 'user', or 'turn N' for something said earlier in this chat. URLs are checked against the pages read in this chat."),
        reason: z.string().optional().describe("For a decision: why it was made."),
      },
      implementation: async ({ kind, text, source, reason }) =>
        withLedger(slug, (l) => {
          if (kind === "finding") return addFinding(l, text, source, "model", chatId);
          if (kind === "decision") return addDecision(l, text, reason, "model", chatId);
          return addQuestion(l, text, "model", chatId);
        }),
    }),
    tool({
      name: "ledger_resolve",
      description: "Close an open question with its answer. Optionally record the answer as a finding, with its source.",
      parameters: {
        question: z.string().describe("The question id, such as Q2."),
        answer: z.string(),
        record_as_finding: z.boolean().optional(),
        source: z.string().optional(),
      },
      implementation: async ({ question, answer, record_as_finding, source }) =>
        withLedger(slug, (l) => resolveQuestion(l, question, answer, record_as_finding ? { source } : null, "model", chatId)),
    }),
  ];
}

export function recallTool(chatId: string): Tool {
  return tool({
    name: "recall",
    description:
      "Search this chat's earlier messages, including parts that compression has summarized away, and get back the original wording with turn numbers. Use it when the user refers to something said earlier that you cannot see in full, or when you need an exact figure, name or quote from earlier in this chat.",
    parameters: {
      query: z.string().describe("Words likely to appear in the passage you want."),
      max_results: z.number().int().min(1).max(8).optional(),
    },
    implementation: async ({ query, max_results }) => formatResults(search(chatId, query, { limit: max_results ?? 4 })),
  });
}

function recallEnabled(ctl: ToolsProviderController): boolean {
  try {
    return ctl.getPluginConfig(configSchematics).get("recall");
  } catch {
    return true;
  }
}

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  const chatId = chatIdOf(ctl);
  if (!chatId) return [];
  const state = loadChatState(chatId);
  const tools: Tool[] = [];
  if (state?.project && loadLedger(state.project)) tools.push(...ledgerTools(chatId, state.project));
  if (!state?.recallOff && recallEnabled(ctl) && hasIndex(chatId)) tools.push(recallTool(chatId));
  return tools;
}

// The model's tools. Offered only in a chat bound to a project; everywhere else the list is
// empty, so ordinary chats pay nothing. Each result is a few lines: what changed, where the
// plan stands, and what to do next. The full state is never repeated in a tool result.

import { tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { chatIdOf, loadChatState } from "./chat";
import { configSchematics } from "./config";
import { formatResults, hasIndex, search } from "./recall";
import { currentStep, Ledger, loadLedger, progress, readEvents, renderMarkdown, stepLabel } from "./ledger";
import { addDecision, addFinding, addQuestion, addSteps, nextLine, OpError, resolveQuestion, updateStep } from "./ops";
import { clip } from "./text";
import { absorbPendingEdit } from "./checks";

const STATUS = ["todo", "doing", "done", "blocked"] as const;
const VIEW_MAX_CHARS = 14000;

function nowLine(l: Ledger): string {
  const p = progress(l);
  const cur = currentStep(l);
  if (!p.total) return "now: no plan yet";
  return `now: ${cur ? clip(stepLabel(cur), 160) : "no step in progress"} | ${p.done} of ${p.total} done`;
}

function result(l: Ledger, changed: string, next?: string): string {
  return [`status: ok`, `changed: ${changed}`, nowLine(l), next ?? nextLine(l)].join("\n");
}

const RUN_AHEAD =
  "next: The plan was finished earlier in this reply, so these steps go beyond what was asked. Unless the user's latest message asked for them, give the user the result now, name these steps as what could come next, and stop.";

const NEXT_STEP_SAME_REPLY =
  "next: A step was finished earlier in this reply. If the user's latest message asked only for that step, give the user its result now and stop; start this step when they ask for it. An automatic note from context management telling you to continue is not a new request from the user.";

// What happened to the plan during this reply, from the event log: the steps the model added
// after finishing the whole plan, and whether it finished any step at all. Local models tend
// to keep going through a list the user means to take one message at a time, especially
// after context compression tells them to "continue the task"; the tool result is the one
// place to say so at the moment it happens.
export function replyProgress(l: Ledger, turnAt: string | undefined): { beyondPlan: Set<number>; finished: Set<number> } {
  const beyondPlan = new Set<number>();
  const finished = new Set<number>();
  if (!turnAt) return { beyondPlan, finished };
  let planDone = false;
  const plan = new Map<string, string>();
  for (const e of readEvents(l.name, 0)) {
    if (e.op === "step_added" && e.id) plan.set(e.id, "todo");
    if (e.op === "step_updated" && e.id) {
      const m = /now (todo|doing|done|blocked)/.exec(e.text ?? "");
      if (m) plan.set(e.id, m[1]);
    }
    if (e.ts < turnAt) continue;
    if (e.op === "step_updated" && e.id && /now done/.test(e.text ?? "")) {
      finished.add(Number(e.id));
      if ([...plan.values()].every((s) => s === "done")) planDone = true;
    } else if (planDone && e.op === "step_added" && e.id) beyondPlan.add(Number(e.id));
  }
  return { beyondPlan, finished };
}

function failure(message: string, next: string): string {
  return [`status: error`, `reason: ${message}`, `next: ${next}`].join("\n");
}

// Every tool reloads the ledger when called, so changes made by commands or other chats are seen.
function withLedger(slug: string, fn: (l: Ledger) => string): string {
  const l = loadLedger(slug);
  if (!l) return failure("This chat's project no longer exists.", "Tell the user; they can start or open a project with /tl.");
  try {
    // The user may have saved an edit to ledger.md since the last message: merge it first.
    absorbPendingEdit(l);
    return result(l, fn(l));
  } catch (e) {
    if (e instanceof OpError) return failure(e.message, "Fix the call and try again, or continue without it.");
    return failure("The ledger could not be written.", "Continue without recording, and tell the user if it keeps failing.");
  }
}

function viewSection(l: Ledger, section: string): string {
  const md = renderMarkdown(l);
  if (section === "all") return md;
  const heading: Record<string, string> = {
    goal: "## Goal", constraints: "## Constraints", plan: "## Plan", findings: "## Findings", decisions: "## Decisions", questions: "## Open questions",
  };
  const start = md.indexOf(heading[section]);
  if (start < 0) return md;
  const rest = md.slice(start);
  const end = rest.indexOf("\n## ", 3);
  return end > 0 ? rest.slice(0, end) : rest;
}

export function ledgerTools(chatId: string, slug: string): Tool[] {
  return [
    tool({
      name: "ledger_view",
      description:
        "Show this chat's project ledger: goal, constraints, plan, findings with sources, decisions and questions. Use it when the throughLine state in the user's message is not enough.",
      parameters: { section: z.enum(["all", "goal", "constraints", "plan", "findings", "decisions", "questions"]).optional() },
      implementation: async ({ section }) => {
        const l = loadLedger(slug);
        if (!l) return failure("This chat's project no longer exists.", "Tell the user.");
        let text = viewSection(l, section ?? "all");
        if (text.length > VIEW_MAX_CHARS) text = text.slice(0, VIEW_MAX_CHARS) + "\n(cut here; view one section at a time for the rest)";
        return `status: ok\n${text}\n${nextLine(l)}`;
      },
    }),
    tool({
      name: "ledger_plan",
      description:
        "Add steps to the project's plan, or change a step's status, note or wording. Plan only what the goal needs; add steps beyond it only if the user asks. Keep exactly one step 'doing' while you work, and mark it 'done' when finished.",
      parameters: {
        action: z.enum(["add", "update"]),
        steps: z.array(z.string()).optional().describe("For add: one or more step texts, in order."),
        step: z.number().int().optional().describe("For update: the step number."),
        status: z.enum(STATUS).optional(),
        note: z.string().optional().describe("A short note on the step, such as what is left or why it is blocked."),
        text: z.string().optional().describe("For update: new wording for the step."),
      },
      implementation: async ({ action, steps, step, status, note, text }) => {
        let warning: string | null = null;
        const out = withLedger(slug, (l) => {
          const turnAt = loadChatState(chatId)?.turnAt;
          if (action === "add") {
            const changed = addSteps(l, steps?.length ? steps : text ? [text] : [], status ?? "todo", "model", chatId);
            if (replyProgress(l, turnAt).beyondPlan.has(l.plan.at(-1)!.id)) warning = RUN_AHEAD;
            return changed;
          }
          if (step === undefined) throw new OpError("For update, give the step number.");
          const changed = updateStep(l, step, { status, note, text }, "model", chatId);
          if (status === "doing") {
            const p = replyProgress(l, turnAt);
            if (p.beyondPlan.has(step)) warning = RUN_AHEAD;
            else if ([...p.finished].some((id) => id !== step)) warning = NEXT_STEP_SAME_REPLY;
          }
          return changed;
        });
        return warning ? out.replace(/\nnext: .*$/, "\n" + warning) : out;
      },
    }),
    tool({
      name: "ledger_record",
      description:
        "Record something established for the project: a finding (a fact, with the source it came from), a decision (with the reason), or an open question that blocks progress. Project facts only; facts about the user belong in memory.",
      parameters: {
        kind: z.enum(["finding", "decision", "question"]),
        text: z.string(),
        source: z.string().optional().describe("For a finding: the URL of the page you read it on, a file path, 'user', or 'turn N' for something said earlier in this chat. URLs are checked against the pages read in this chat."),
        reason: z.string().optional().describe("For a decision: why it was made."),
      },
      implementation: async ({ kind, text, source, reason }) =>
        withLedger(slug, (l) => {
          if (kind === "finding") return addFinding(l, text, source, "model", chatId);
          if (kind === "decision") return addDecision(l, text, reason, "model", chatId);
          return addQuestion(l, text, "model", chatId);
        }),
    }),
    tool({
      name: "ledger_resolve",
      description: "Close an open question with its answer. Optionally record the answer as a finding, with its source.",
      parameters: {
        question: z.string().describe("The question id, such as Q2."),
        answer: z.string(),
        record_as_finding: z.boolean().optional(),
        source: z.string().optional(),
      },
      implementation: async ({ question, answer, record_as_finding, source }) =>
        withLedger(slug, (l) => resolveQuestion(l, question, answer, record_as_finding ? { source } : null, "model", chatId)),
    }),
  ];
}

export function recallTool(chatId: string): Tool {
  return tool({
    name: "recall",
    description:
      "Search this chat's earlier messages, including parts that compression has summarized away, and get back the original wording with turn numbers. Use it when the user refers to something said earlier that you cannot see in full, or when you need an exact figure, name or quote from earlier in this chat.",
    parameters: {
      query: z.string().describe("Words likely to appear in the passage you want."),
      max_results: z.number().int().min(1).max(8).optional(),
    },
    implementation: async ({ query, max_results }) => formatResults(search(chatId, query, { limit: max_results ?? 4 })),
  });
}

function recallEnabled(ctl: ToolsProviderController): boolean {
  try {
    return ctl.getPluginConfig(configSchematics).get("recall");
  } catch {
    return true;
  }
}

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  const chatId = chatIdOf(ctl);
  if (!chatId) return [];
  const state = loadChatState(chatId);
  const tools: Tool[] = [];
  if (state?.project && loadLedger(state.project)) tools.push(...ledgerTools(chatId, state.project));
  if (!state?.recallOff && recallEnabled(ctl) && hasIndex(chatId)) tools.push(recallTool(chatId));
  return tools;
}