// Tracks one research task at a time: an unbroken run of searches and page reads, its
// central goal, and what the run has cost. Every tool response repeats the central goal
// in a "task:" line (and the current narrower goal in a "step:" line), so the task stays
// in the newest part of the context even after older turns are compressed.

import { terms } from "./focus";

export interface Research {
  id: number;
  // The first goal stated in this run of calls: the central task.
  goal: string;
  key: string;
  inferred: boolean;
  // The narrower goal of the current call, when it differs from the central task.
  step: string | null;
  searches: number;
  pages: number;
  chars: number;
  lastAt: number;
  softNoted: boolean;
  exhausted: boolean;
  // Every goal stated in this run, normalized, and how many calls were refused.
  seen: string[];
  refusals: number;
}

// A research task is one unbroken run of tool calls. Within a reply a model pauses for
// seconds between calls; a gap this long means the reply ended and a new request came.
// The budget belongs to the run, so narrowing or rewording the goal does not reset it.
const BURST_GAP_MS = 90_000;
// A shorter pause also ends the run when the next call states a goal not seen in it: the model
// wrote its answer and the user asked something new, quickly. Within one reply, calls that
// follow each other closely keep their budget whatever goal they name.
const TURN_GAP_MS = 25_000;

const normalize = (s: string) =>
  s
    .toLowerCase()
    .replace(/[^\p{L}\p{N}]+/gu, " ")
    .trim();

// Goals sharing most of their content words describe the same thing.
export function sameTask(a: string, b: string): boolean {
  const ta = new Set(terms(a));
  const tb = new Set(terms(b));
  if (!ta.size || !tb.size) return false;
  let shared = 0;
  for (const t of ta) if (tb.has(t)) shared++;
  return shared / Math.min(ta.size, tb.size) >= 0.6;
}

export class ResearchTracker {
  current: Research | null = null;
  private nextId = 1;
  private hardPages: number;

  constructor(
    private budgetTokens: number,
    private enforce = true,
    private log: (msg: string) => void = () => {},
    private softPages = 6
  ) {
    this.hardPages = Math.ceil(softPages * 1.7);
  }

  private start(goal: string, inferred: boolean, now: number): Research {
    const r: Research = {
      id: this.nextId++,
      goal: goal.slice(0, 300),
      key: normalize(goal),
      inferred,
      step: null,
      searches: 0,
      pages: 0,
      chars: 0,
      lastAt: now,
      softNoted: false,
      exhausted: false,
      seen: [normalize(goal)],
      refusals: 0,
    };
    this.log(`research task ${r.id} started`);
    return r;
  }

  // The user sent a new message, so the next call starts a new task with a fresh budget. Only
  // the LM Studio plugin can tell this; the MCP server relies on the pauses above.
  newTurn(): void {
    if (this.current) this.log(`research task ${this.current.id} ended with a new message`);
    this.current = null;
  }

  resolve(goal: string | undefined, fallbackQuery?: string): Research | null {
    const now = Date.now();
    const g = (goal ?? "").trim().slice(0, 300);
    if (this.current) {
      const idle = now - this.current.lastAt;
      const newGoal = !!g && !this.current.seen.includes(normalize(g));
      if (idle > BURST_GAP_MS || (idle > TURN_GAP_MS && newGoal)) this.current = null;
    }
    if (!this.current) {
      if (g) this.current = this.start(g, false, now);
      else if ((fallbackQuery ?? "").trim()) this.current = this.start((fallbackQuery ?? "").trim(), true, now);
      else return null;
      return this.current;
    }
    const r = this.current;
    if (g) {
      const key = normalize(g);
      if (r.inferred) {
        // The first search had no goal and was named after its query; the first stated
        // goal becomes the central task.
        r.goal = g;
        r.key = normalize(g);
        r.inferred = false;
        r.step = null;
      } else if (normalize(g) === r.key || sameTask(r.goal, g)) {
        r.step = null;
      } else {
        r.step = g;
      }
      if (!r.seen.includes(key)) r.seen = [...r.seen, key].slice(-20);
    }
    return r;
  }

  // The goal a page read should focus on: the current step if there is one.
  focusGoal(r: Research | null): string | undefined {
    if (!r || r.inferred) return undefined;
    return r.step ?? r.goal;
  }

  record(r: Research, kind: "search" | "page", chars: number): void {
    if (kind === "search") r.searches++;
    else r.pages++;
    r.chars += chars;
    r.lastAt = Date.now();
  }

  tokens(r: Research): number {
    return Math.round(r.chars / 4 / 100) * 100;
  }

  level(r: Research): "ok" | "soft" | "hard" {
    const t = r.chars / 4;
    if (t >= this.budgetTokens * 1.75 || r.pages >= this.hardPages) return "hard";
    if (t >= this.budgetTokens || r.pages >= this.softPages) return "soft";
    return "ok";
  }

  // Logs counts only, never goals or URLs.
  noteLevel(r: Research, level: "ok" | "soft" | "hard"): void {
    const counts = `${r.searches} search(es), ${r.pages} read(s), about ${this.tokens(r)} tokens`;
    if (level === "hard" && !r.exhausted) {
      r.exhausted = true;
      this.log(`research task ${r.id}: budget used up after ${counts}${this.enforce ? "; further calls in this run are refused" : ""}`);
    } else if (level === "soft" && !r.softNoted) {
      r.softNoted = true;
      this.log(`research task ${r.id}: wrap-up point reached after ${counts}`);
    }
  }

  // With enforcement on, once the stop instruction has been delivered, the rest of the
  // run gets no more searches or reads: a model can ignore advice, not a refusal.
  isBlocked(r: Research | null): boolean {
    return this.enforce && !!r && r.exhausted;
  }

  // Refusals keep the run alive, so a model that keeps calling stays refused until it answers.
  refuse(r: Research, kind: "search" | "page"): string {
    r.lastAt = Date.now();
    r.refusals++;
    this.log(`research task ${r.id}: refused a ${kind === "search" ? "search" : "page read"} after the budget ran out (refusal ${r.refusals})`);
    return [
      "status: budget_exhausted",
      `task: ${r.goal}`,
      `budget: ${r.searches} search(es) and ${r.pages} page read(s) for this task, about ${this.tokens(r).toLocaleString("en-US")} tokens of tool output`,
      r.refusals > 1
        ? `next: Refused again (${r.refusals} times now). Every search and page read in this reply will be refused the same way, whatever its query, URL or goal. Stop calling web_search and read_webpage, and write your answer now.`
        : "next: This call was not run because the research budget for this reply is used up, and further calls in this reply will be refused too. Answer the user now with what you already have and name any gaps. If work remains, say what is left; the user's next message gets a new budget.",
    ].join("\n");
  }
}

const ADVANCING = /^status: (ok|no_results|parser_failed)\b/;

// Adds the task and budget lines to a tool response and, once the budget is spent,
// replaces its "next:" line with an instruction to wrap up.
export function annotate(output: string, r: Research | null, tracker: ResearchTracker, kind: "search" | "page"): string {
  if (!r) return output;
  const lines = output.split("\n");
  const header = [`task: ${r.goal}${r.inferred ? " (taken from the first search; pass goal to name the task)" : ""}`];
  if (r.step) header.push(`step: ${r.step}`);
  lines.splice(1, 0, ...header);
  tracker.record(r, kind, output.length + header.join("\n").length);

  const level = tracker.level(r);
  tracker.noteLevel(r, level);
  const budgetLine = `budget: ${r.searches} search(es) and ${r.pages} page read(s) so far for this task, about ${tracker
    .tokens(r)
    .toLocaleString("en-US")} tokens of tool output`;
  let nextIndex = -1;
  for (let i = lines.length - 1; i >= 0; i--) {
    if (lines[i].startsWith("next:")) {
      nextIndex = i;
      break;
    }
  }
  const insertAt = nextIndex >= 0 ? nextIndex : lines.length;
  lines.splice(insertAt, 0, budgetLine);
  if (nextIndex >= 0) nextIndex++;

  if (ADVANCING.test(lines[0]) && level !== "ok") {
    const wrap =
      level === "hard"
        ? "next: The research budget for this reply is used up, and further searches and page reads in this reply will be refused. Answer now with what you have, naming any gaps. If work remains, say what is left; the user's next message gets a new budget."
        : "next: You have already gathered a lot for this task. Answer now with what you have, unless one specific gap remains; if it does, close it with one more search or read at most.";
    if (nextIndex >= 0) lines[nextIndex] = wrap;
    else lines.push(wrap);
  }
  return lines.join("\n");
}

// Tracks one research task at a time: an unbroken run of searches and page reads, its
// central goal, and what the run has cost. Every tool response repeats the central goal
// in a "task:" line (and the current narrower goal in a "step:" line), so the task stays
// in the newest part of the context even after older turns are compressed.

import { terms } from "./focus";

export interface Research {
  id: number;
  // The first goal stated in this run of calls: the central task.
  goal: string;
  key: string;
  inferred: boolean;
  // The narrower goal of the current call, when it differs from the central task.
  step: string | null;
  searches: number;
  pages: number;
  chars: number;
  lastAt: number;
  softNoted: boolean;
  exhausted: boolean;
  // Every goal stated in this run, normalized, and how many calls were refused.
  seen: string[];
  refusals: number;
}

// A research task is one unbroken run of tool calls. Within a reply a model pauses for
// seconds between calls; a gap this long means the reply ended and a new request came.
// The budget belongs to the run, so narrowing or rewording the goal does not reset it.
const BURST_GAP_MS = 90_000;
// A shorter pause also ends the run when the next call states a goal not seen in it: the model
// wrote its answer and the user asked something new, quickly. Within one reply, calls that
// follow each other closely keep their budget whatever goal they name.
const TURN_GAP_MS = 25_000;

const normalize = (s: string) =>
  s
    .toLowerCase()
    .replace(/[^\p{L}\p{N}]+/gu, " ")
    .trim();

// Goals sharing most of their content words describe the same thing.
export function sameTask(a: string, b: string): boolean {
  const ta = new Set(terms(a));
  const tb = new Set(terms(b));
  if (!ta.size || !tb.size) return false;
  let shared = 0;
  for (const t of ta) if (tb.has(t)) shared++;
  return shared / Math.min(ta.size, tb.size) >= 0.6;
}

export class ResearchTracker {
  current: Research | null = null;
  private nextId = 1;
  private hardPages: number;

  constructor(
    private budgetTokens: number,
    private enforce = true,
    private log: (msg: string) => void = () => {},
    private softPages = 6
  ) {
    this.hardPages = Math.ceil(softPages * 1.7);
  }

  private start(goal: string, inferred: boolean, now: number): Research {
    const r: Research = {
      id: this.nextId++,
      goal: goal.slice(0, 300),
      key: normalize(goal),
      inferred,
      step: null,
      searches: 0,
      pages: 0,
      chars: 0,
      lastAt: now,
      softNoted: false,
      exhausted: false,
      seen: [normalize(goal)],
      refusals: 0,
    };
    this.log(`research task ${r.id} started`);
    return r;
  }

  // The user sent a new message, so the next call starts a new task with a fresh budget. Only
  // the LM Studio plugin can tell this; the MCP server relies on the pauses above.
  newTurn(): void {
    if (this.current) this.log(`research task ${this.current.id} ended with a new message`);
    this.current = null;
  }

  resolve(goal: string | undefined, fallbackQuery?: string): Research | null {
    const now = Date.now();
    const g = (goal ?? "").trim().slice(0, 300);
    if (this.current) {
      const idle = now - this.current.lastAt;
      const newGoal = !!g && !this.current.seen.includes(normalize(g));
      if (idle > BURST_GAP_MS || (idle > TURN_GAP_MS && newGoal)) this.current = null;
    }
    if (!this.current) {
      if (g) this.current = this.start(g, false, now);
      else if ((fallbackQuery ?? "").trim()) this.current = this.start((fallbackQuery ?? "").trim(), true, now);
      else return null;
      return this.current;
    }
    const r = this.current;
    if (g) {
      const key = normalize(g);
      if (r.inferred) {
        // The first search had no goal and was named after its query; the first stated
        // goal becomes the central task.
        r.goal = g;
        r.key = normalize(g);
        r.inferred = false;
        r.step = null;
      } else if (normalize(g) === r.key || sameTask(r.goal, g)) {
        r.step = null;
      } else {
        r.step = g;
      }
      if (!r.seen.includes(key)) r.seen = [...r.seen, key].slice(-20);
    }
    return r;
  }

  // The goal a page read should focus on: the current step if there is one.
  focusGoal(r: Research | null): string | undefined {
    if (!r || r.inferred) return undefined;
    return r.step ?? r.goal;
  }

  record(r: Research, kind: "search" | "page", chars: number): void {
    if (kind === "search") r.searches++;
    else r.pages++;
    r.chars += chars;
    r.lastAt = Date.now();
  }

  tokens(r: Research): number {
    return Math.round(r.chars / 4 / 100) * 100;
  }

  level(r: Research): "ok" | "soft" | "hard" {
    const t = r.chars / 4;
    if (t >= this.budgetTokens * 1.75 || r.pages >= this.hardPages) return "hard";
    if (t >= this.budgetTokens || r.pages >= this.softPages) return "soft";
    return "ok";
  }

  // Logs counts only, never goals or URLs.
  noteLevel(r: Research, level: "ok" | "soft" | "hard"): void {
    const counts = `${r.searches} search(es), ${r.pages} read(s), about ${this.tokens(r)} tokens`;
    if (level === "hard" && !r.exhausted) {
      r.exhausted = true;
      this.log(`research task ${r.id}: budget used up after ${counts}${this.enforce ? "; further calls in this run are refused" : ""}`);
    } else if (level === "soft" && !r.softNoted) {
      r.softNoted = true;
      this.log(`research task ${r.id}: wrap-up point reached after ${counts}`);
    }
  }

  // With enforcement on, once the stop instruction has been delivered, the rest of the
  // run gets no more searches or reads: a model can ignore advice, not a refusal.
  isBlocked(r: Research | null): boolean {
    return this.enforce && !!r && r.exhausted;
  }

  // Refusals keep the run alive, so a model that keeps calling stays refused until it answers.
  refuse(r: Research, kind: "search" | "page"): string {
    r.lastAt = Date.now();
    r.refusals++;
    this.log(`research task ${r.id}: refused a ${kind === "search" ? "search" : "page read"} after the budget ran out (refusal ${r.refusals})`);
    return [
      "status: budget_exhausted",
      `task: ${r.goal}`,
      `budget: ${r.searches} search(es) and ${r.pages} page read(s) for this task, about ${this.tokens(r).toLocaleString("en-US")} tokens of tool output`,
      r.refusals > 1
        ? `next: Refused again (${r.refusals} times now). Every search and page read in this reply will be refused the same way, whatever its query, URL or goal. Stop calling web_search and read_webpage, and write your answer now.`
        : "next: This call was not run because the research budget for this reply is used up, and further calls in this reply will be refused too. Answer the user now with what you already have and name any gaps. If work remains, say what is left; the user's next message gets a new budget.",
    ].join("\n");
  }
}

const ADVANCING = /^status: (ok|no_results|parser_failed)\b/;

// Adds the task and budget lines to a tool response and, once the budget is spent,
// replaces its "next:" line with an instruction to wrap up.
export function annotate(output: string, r: Research | null, tracker: ResearchTracker, kind: "search" | "page"): string {
  if (!r) return output;
  const lines = output.split("\n");
  const header = [`task: ${r.goal}${r.inferred ? " (taken from the first search; pass goal to name the task)" : ""}`];
  if (r.step) header.push(`step: ${r.step}`);
  lines.splice(1, 0, ...header);
  tracker.record(r, kind, output.length + header.join("\n").length);

  const level = tracker.level(r);
  tracker.noteLevel(r, level);
  const budgetLine = `budget: ${r.searches} search(es) and ${r.pages} page read(s) so far for this task, about ${tracker
    .tokens(r)
    .toLocaleString("en-US")} tokens of tool output`;
  let nextIndex = -1;
  for (let i = lines.length - 1; i >= 0; i--) {
    if (lines[i].startsWith("next:")) {
      nextIndex = i;
      break;
    }
  }
  const insertAt = nextIndex >= 0 ? nextIndex : lines.length;
  lines.splice(insertAt, 0, budgetLine);
  if (nextIndex >= 0) nextIndex++;

  if (ADVANCING.test(lines[0]) && level !== "ok") {
    const wrap =
      level === "hard"
        ? "next: The research budget for this reply is used up, and further searches and page reads in this reply will be refused. Answer now with what you have, naming any gaps. If work remains, say what is left; the user's next message gets a new budget."
        : "next: You have already gathered a lot for this task. Answer now with what you have, unless one specific gap remains; if it does, close it with one more search or read at most.";
    if (nextIndex >= 0) lines[nextIndex] = wrap;
    else lines.push(wrap);
  }
  return lines.join("\n");
}