import { describe, it, expect, vi } from "vitest";
import { PromptGuard, registerPromptGuard, modelKeyOf, noteOversizedRequest } from "./guard.js";
import { DEFAULT_HEADROOM_TOKENS, type PromptcapSettings } from "./limits.js";
import type { AgentMessage } from "./fold.js";

const settings = (over: Partial<PromptcapSettings> = {}): PromptcapSettings => ({
  enabled: true,
  perModel: {},
  ...over,
});

const conversation = (n: number, size: number): AgentMessage[] => {
  const messages: AgentMessage[] = [{ role: "user", content: [{ type: "text", text: "go" }] }];
  for (let i = 0; i < n; i++) {
    messages.push({ role: "assistant", content: [{ type: "toolCall", id: `t${i}`, name: "read", arguments: { path: `/f${i}` } }] });
    messages.push({ role: "toolResult", toolCallId: `t${i}`, toolName: "read", content: [{ type: "text", text: "o".repeat(size) }], isError: false });
  }
  return messages;
};

const ctx = (window?: number) => ({
  model: { provider: "anthropic", id: "claude-opus-4-8" },
  getSystemPrompt: () => "system",
  getContextUsage: () => (window ? { tokens: null, contextWindow: window, percent: null } : undefined),
});

describe("PromptGuard", () => {
  it("leaves the conversation alone when folding is disabled", () => {
    const guard = new PromptGuard({ settings: () => settings({ enabled: false }) });
    const messages = conversation(40, 8000);
    const before = JSON.parse(JSON.stringify(messages));

    expect(guard.apply(messages, ctx(200_000), [])).toEqual(before);
    expect(guard.lastTokens).toBeNull();
  });

  it("sizes an image by the model's declared image cost", () => {
    const withImage = (): AgentMessage[] => [
      { role: "user", content: [{ type: "text", text: "go" }] },
      { role: "assistant", content: [{ type: "toolCall", id: "t0", name: "read", arguments: { path: "/a.png" } }] },
      { role: "toolResult", toolCallId: "t0", toolName: "read", content: [{ type: "image", data: "b".repeat(4_000), mimeType: "image/png" }], isError: false },
      { role: "user", content: [{ type: "text", text: "next" }] },
    ];
    const cheap = new PromptGuard({ settings: () => settings() });
    const dear = new PromptGuard({ settings: () => settings({ perModel: { "claude-opus-4-8": { imageTokens: 25_000 } } }) });

    cheap.apply(withImage(), ctx(200_000), []);
    dear.apply(withImage(), ctx(200_000), []);

    expect(dear.lastTokens! - cheap.lastTokens!).toBe(25_000 - 1_600);
  });

  it("asks for the hour-long cache unless the operator already chose", () => {
    const original = process.env.PI_CACHE_RETENTION;
    try {
      delete process.env.PI_CACHE_RETENTION;
      new PromptGuard({ settings: () => settings() }).apply(conversation(2, 100), ctx(200_000), []);
      expect(process.env.PI_CACHE_RETENTION).toBe("long");

      process.env.PI_CACHE_RETENTION = "short";
      new PromptGuard({ settings: () => settings() }).apply(conversation(2, 100), ctx(200_000), []);
      expect(process.env.PI_CACHE_RETENTION).toBe("short");

      delete process.env.PI_CACHE_RETENTION;
      new PromptGuard({ settings: () => settings({ longCacheRetention: false }) }).apply(conversation(2, 100), ctx(200_000), []);
      expect(process.env.PI_CACHE_RETENTION).toBeUndefined();

      // Withdrawing the setting mid-session gives the provider default back.
      delete process.env.PI_CACHE_RETENTION;
      let wanted = true;
      const guard = new PromptGuard({ settings: () => settings({ longCacheRetention: wanted }) });
      guard.apply(conversation(2, 100), ctx(200_000), []);
      expect(process.env.PI_CACHE_RETENTION).toBe("long");
      wanted = false;
      guard.apply(conversation(2, 100), ctx(200_000), []);
      expect(process.env.PI_CACHE_RETENTION).toBeUndefined();
    } finally {
      if (original === undefined) delete process.env.PI_CACHE_RETENTION;
      else process.env.PI_CACHE_RETENTION = original;
    }
  });

  it("names the message a fold rewrote, so its re-bill can be traced", () => {
    const log = vi.fn();
    const guard = new PromptGuard({ settings: () => settings({ maxPromptTokens: 5_000 }), log });

    guard.apply(conversation(40, 8000), ctx(), []);

    const [event, message] = log.mock.calls.find(([e]) => (e as any).rewroteFrom !== undefined)!;
    expect((event as any).rewroteFrom).toBe(1);
    expect((event as any).promoted).toBeGreaterThan(0);
    expect(message).toContain("rewrote");
  });

  it("reports a fold held back as not worth its cache miss", () => {
    const log = vi.fn();
    const guard = new PromptGuard({ settings: () => settings({ maxPromptTokens: 19_000 }), log });
    const messages: AgentMessage[] = [
      { role: "user", content: [{ type: "text", text: "p".repeat(80_000) }] },
      ...conversation(10, 200).slice(1),
    ];

    guard.apply(messages, ctx(), []);

    expect(log.mock.calls.some(([, m]) => String(m).includes("not free enough"))).toBe(true);
    expect(log.mock.calls.some(([e]) => (e as any).rewroteFrom !== undefined)).toBe(false);
  });

  it("folds a conversation that outgrew the declared ceiling", () => {
    const guard = new PromptGuard({ settings: () => settings({ maxPromptTokens: 5_000 }) });
    const messages = conversation(40, 8000);

    guard.apply(messages, ctx(), []);

    expect(guard.lastTokens!).toBeLessThanOrEqual(5_000);
    expect(messages[2].content[0].text).toBe("[omitted: 8000B]");
  });

  it("warns once the conversation cannot be made to fit", () => {
    const notify = vi.fn();
    const guard = new PromptGuard({ settings: () => settings({ maxPromptTokens: 1_000 }), notify });
    // Prose alone overflows, and prose is never folded.
    guard.apply([{ role: "user", content: [{ type: "text", text: "p".repeat(400_000) }] }], ctx(), []);

    expect(notify).toHaveBeenCalledWith(expect.stringContaining("no longer fits"), "warning");
  });

  it("does not warn when folding brought the prompt back under the ceiling", () => {
    const notify = vi.fn();
    const guard = new PromptGuard({ settings: () => settings({ maxPromptTokens: 5_000 }), notify });
    guard.apply(conversation(40, 8000), ctx(), []);

    expect(notify).not.toHaveBeenCalled();
  });

  it("uses the window the host reports when none is declared", () => {
    const guard = new PromptGuard({ settings: () => settings() });
    guard.apply(conversation(2, 100), ctx(1_000_000), []);

    // The headroom above the prose that cannot be folded, which a reported
    // window is enough to unlock: without one the ceiling would stay at the
    // configured default.
    expect(guard.lastCeiling!).toBeGreaterThanOrEqual(DEFAULT_HEADROOM_TOKENS);
    expect(guard.lastCeiling!).toBeLessThan(DEFAULT_HEADROOM_TOKENS + 1_000);
  });

  it("reports what signature-only reasoning blocks weighed against what was charged", () => {
    const log = vi.fn();
    const guard = new PromptGuard({ settings: () => settings(), log });
    const messages = conversation(2, 100);
    // A reasoning block as the subscription gateway leaves them: a signature
    // and no text. Whether the adapter sends it is what the log settles.
    messages.splice(1, 0, {
      role: "assistant",
      content: [{ type: "thinking", thinking: "", thinkingSignature: "s".repeat(4000) }],
    });
    guard.apply(messages, ctx(1_000_000), []);

    guard.calibrate(12_345, "anthropic/claude-opus-4-8");

    const reported = log.mock.calls.find(([event]) => (event as any).danglingSignatures > 0);
    expect(reported).toBeDefined();
    const [event] = reported!;
    expect((event as any).charged).toBe(12_345);
    expect((event as any).danglingSignatures).toBe(1000);
    expect((event as any).predictedWithout).toBe((event as any).predicted - 1000);
  });

  it("says nothing about signatures when there are none to explain", () => {
    const log = vi.fn();
    const guard = new PromptGuard({ settings: () => settings(), log });
    guard.apply(conversation(2, 100), ctx(1_000_000), []);

    guard.calibrate(999, "anthropic/claude-opus-4-8");

    expect(log.mock.calls.some(([event]) => (event as any).danglingSignatures !== undefined)).toBe(false);
  });

  it("keeps calibration on the raw estimate, so the ratio holds where it lands", () => {
    const guard = new PromptGuard({ settings: () => settings() });
    const messages = conversation(2, 4000);

    // The first sizing is uncalibrated, so it is the raw estimate itself.
    guard.apply(messages, ctx(1_000_000), []);
    const raw = guard.lastTokens!;

    // A provider that consistently charges twice what the bytes suggest. Fed
    // its own scaled prediction instead, the ratio would drift to the square
    // root of the truth and understate every prompt by a third.
    guard.calibrate(raw * 2, "anthropic/claude-opus-4-8");
    for (let i = 0; i < 40; i++) {
      guard.apply(messages, ctx(1_000_000), []);
      guard.calibrate(raw * 2, "anthropic/claude-opus-4-8");
    }

    expect(guard.ratioFor("anthropic/claude-opus-4-8")!).toBeGreaterThan(1.9);
    expect(guard.ratioFor("anthropic/claude-opus-4-8")!).toBeLessThan(2.1);
  });

  it("learns the ratio from what the provider charged", () => {
    const guard = new PromptGuard({ settings: () => settings() });
    guard.apply(conversation(2, 4000), ctx(), []);
    const predicted = guard.lastTokens!;

    guard.calibrate(predicted * 2, "anthropic/claude-opus-4-8");

    expect(guard.ratioFor("anthropic/claude-opus-4-8")).toBeCloseTo(2, 3);
  });

  it("ignores a charge from a model other than the one it sized for", () => {
    const guard = new PromptGuard({ settings: () => settings() });
    guard.apply(conversation(2, 4000), ctx(), []);
    guard.calibrate(999_999, "github-copilot/gpt-5.6-sol");

    expect(guard.ratioFor("github-copilot/gpt-5.6-sol")).toBeUndefined();
    expect(guard.ratioFor("anthropic/claude-opus-4-8")).toBeUndefined();
  });

  it("forgets its folds on reset, so a new session starts whole", () => {
    const guard = new PromptGuard({ settings: () => settings({ maxPromptTokens: 1_000 }) });
    guard.apply(conversation(10, 8000), ctx(), []);
    guard.reset();

    const fresh = conversation(2, 100);
    guard.apply(fresh, ctx(), []);

    expect(fresh[2].content[0].text).toBe("o".repeat(100));
  });
});

describe("registerPromptGuard", () => {
  const harness = () => {
    const handlers = new Map<string, Function>();
    const pi = {
      on: (event: string, handler: Function) => handlers.set(event, handler),
      getAllTools: () => [{ name: "read", description: "d", parameters: {} }],
      getActiveTools: () => ["read"],
    } as any;
    return { handlers, pi };
  };

  it("returns folded messages from the context event", () => {
    const { handlers, pi } = harness();
    registerPromptGuard(pi, new PromptGuard({ settings: () => settings({ maxPromptTokens: 5_000 }) }));

    const messages = conversation(40, 8000);
    const out = handlers.get("context")!({ messages }, ctx());

    expect(out.messages).toBe(messages);
    expect(out.messages[2].content[0].text).toMatch(/^\[omitted: /);
  });

  it("ignores a context event without a message array", () => {
    const { handlers, pi } = harness();
    registerPromptGuard(pi, new PromptGuard({ settings: () => settings() }));

    expect(handlers.get("context")!({}, ctx())).toBeUndefined();
  });

  it("calibrates from the whole prompt, cache included", () => {
    const { handlers, pi } = harness();
    const guard = new PromptGuard({ settings: () => settings() });
    registerPromptGuard(pi, guard);

    handlers.get("context")!({ messages: conversation(2, 4000) }, ctx());
    const predicted = guard.lastTokens!;
    handlers.get("turn_end")!({ message: { usage: { input: 1, cacheRead: predicted * 2 - 2, cacheWrite: 1, output: 5 } } }, ctx());

    expect(guard.ratioFor("anthropic/claude-opus-4-8")).toBeCloseTo(2, 3);
  });
});

describe("a request refused for its size", () => {
  const captures = (n: number, size: number): AgentMessage[] => {
    const messages: AgentMessage[] = [{ role: "user", content: [{ type: "text", text: "go" }] }];
    for (let i = 0; i < n; i++) {
      messages.push({ role: "assistant", content: [{ type: "toolCall", id: `t${i}`, name: "read", arguments: { path: `/s${i}.png` } }] });
      messages.push({ role: "toolResult", toolCallId: `t${i}`, toolName: "read", content: [{ type: "image", data: "i".repeat(size), mimeType: "image/png" }], isError: false });
    }
    return messages;
  };
  const imaged = (sessionManager: object) => ({ ...ctx(200_000), sessionManager });

  it("takes every image off the next attempt, not just the oldest", () => {
    const sessionManager = {};
    const guard = new PromptGuard({ settings: () => settings() });
    guard.apply(captures(2, 1000), imaged(sessionManager), []);

    expect(noteOversizedRequest({ sessionManager })).toBe(true);

    // Two captures are nowhere near the payload ceiling, so nothing but the
    // refusal itself can be what empties them.
    const retried = captures(2, 1000);
    guard.apply(retried, imaged(sessionManager), []);
    expect(retried[2].content[0]).toEqual({ type: "text", text: "[omitted: 1000B]" });
    expect(retried[4].content[0]).toEqual({ type: "text", text: "[omitted: 1000B]" });
  });

  it("promises nothing once an attempt found no image left to take", () => {
    const sessionManager = {};
    const guard = new PromptGuard({ settings: () => settings() });
    const messages = conversation(2, 4000);
    guard.apply(messages, imaged(sessionManager), []);

    expect(noteOversizedRequest({ sessionManager })).toBe(true);
    guard.apply(conversation(2, 4000), imaged(sessionManager), []);
    expect(noteOversizedRequest({ sessionManager })).toBe(false);
  });

  it("forgets the refusal once the provider answers a turn", () => {
    const sessionManager = {};
    const guard = new PromptGuard({ settings: () => settings() });
    guard.apply(conversation(2, 4000), imaged(sessionManager), []);
    noteOversizedRequest({ sessionManager });
    guard.apply(conversation(2, 4000), imaged(sessionManager), []);
    expect(noteOversizedRequest({ sessionManager })).toBe(false);

    guard.calibrate(1000, "anthropic/claude-opus-4-8");

    expect(noteOversizedRequest({ sessionManager })).toBe(true);
  });

  // A turn ends on a refusal too, carrying no charge and carrying the very
  // refusal that would be forgotten here.
  it("does not forget the refusal when the failed turn itself ends", () => {
    const handlers = new Map<string, Function>();
    const sessionManager = {};
    const guard = new PromptGuard({ settings: () => settings() });
    registerPromptGuard({ on: (event: string, handler: Function) => handlers.set(event, handler), getAllTools: () => [], getActiveTools: () => [] } as any, guard);
    guard.apply(conversation(2, 4000), imaged(sessionManager), []);
    noteOversizedRequest({ sessionManager });
    guard.apply(conversation(2, 4000), imaged(sessionManager), []);

    handlers.get("turn_end")!({ message: { stopReason: "error", usage: { input: 0, cacheRead: 0, cacheWrite: 0, output: 0 } } }, ctx());

    expect(noteOversizedRequest({ sessionManager })).toBe(false);
  });

  it("promises nothing while folding is switched off", () => {
    const sessionManager = {};
    const guard = new PromptGuard({ settings: () => settings({ enabled: false }) });
    guard.apply(captures(2, 1000), imaged(sessionManager), []);

    expect(noteOversizedRequest({ sessionManager })).toBe(false);
  });

  it("reaches the guard that built the prompt and no other", () => {
    const mine = {};
    const theirs = {};
    const guard = new PromptGuard({ settings: () => settings() });
    guard.apply(captures(2, 1000), imaged(mine), []);

    expect(noteOversizedRequest({ sessionManager: theirs })).toBe(false);
    expect(noteOversizedRequest({})).toBe(false);
    expect(noteOversizedRequest({ sessionManager: mine })).toBe(true);
  });
});

describe("modelKeyOf", () => {
  it("prefers the full spec and falls back to the bare id", () => {
    expect(modelKeyOf({ model: { provider: "p", id: "m" } })).toBe("p/m");
    expect(modelKeyOf({ model: { id: "m" } })).toBe("m");
    expect(modelKeyOf({})).toBe("");
  });
});