import type { AgentMessage } from "@earendil-works/pi-agent-core";
import { createAssistantMessageEventStream, normalizeContext } from "@earendil-works/pi-ai";
import { describe, expect, it } from "vitest";
import {
  describeToolActivity,
  estimateAgentContextTokens,
  maxInputTokensForWindow,
  parseContextOverflowDetails,
  pruneContextToWindow,
  raceSignal,
  streamWithContextOverflowRecovery,
} from "./pi-runtime.js";

describe("describeToolActivity", () => {
  it("summarizes builtin tools with their most informative argument", () => {
    expect(describeToolActivity("shell", { command: "pnpm test --filter web" })).toBe(
      "Running: pnpm test --filter web",
    );
    expect(describeToolActivity("read_file", { path: "notes/plan.md" })).toBe(
      "Reading notes/plan.md",
    );
    expect(describeToolActivity("write_file", { path: "out.csv", content: "…" })).toBe(
      "Writing out.csv",
    );
    expect(describeToolActivity("render_plot", { spec: {} })).toBe("Rendering a chart");
    expect(describeToolActivity("add_mcp_server", { name: "Linear" })).toBe(
      "Connecting MCP server: Linear",
    );
    expect(describeToolActivity("run_subagent", { name: "scout", task: "…" })).toBe(
      "Delegating to helper: scout",
    );
    expect(describeToolActivity("create_space", { name: "Customer support" })).toBe(
      "Creating space: Customer support",
    );
    expect(describeToolActivity("web_search", { query: "rakazo agents" })).toBe(
      "Searching the web: rakazo agents",
    );
    expect(describeToolActivity("web_fetch", { url: "https://example.com" })).toBe(
      "Reading page: https://example.com/",
    );
  });

  it("names MCP server and remote tool", () => {
    expect(describeToolActivity("mcp__brex__list_expenses", {})).toBe("Using brex: list_expenses");
    expect(describeToolActivity("mcp__demo-oauth__greet", {})).toBe("Using demo-oauth: greet");
  });

  it("truncates long details and collapses whitespace", () => {
    const long = `x${"y".repeat(200)}`;
    const line = describeToolActivity("shell", { command: `a\n\t${long}` });
    expect(line.length).toBeLessThanOrEqual("Running: ".length + 91);
    expect(line).toContain("…");
    expect(line).not.toContain("\n");
    expect(line).toMatch(/^Running: a x/);
  });

  it("redacts credentials from activity details", () => {
    const token = "fake-token";
    const line = describeToolActivity("shell", {
      command: `curl -H 'Authorization: Bearer ${token}' https://example.test?api_key=fake-key password=fake-password`,
    });

    expect(line).toContain("Bearer [redacted]");
    expect(line).toContain("api_key=[redacted]");
    expect(line).not.toContain(token);
    expect(line).not.toContain("fake-key");
    expect(line).not.toContain("fake-password");
  });

  it("strips signed-URL query and fragment from web_fetch activity", () => {
    const line = describeToolActivity("web_fetch", {
      url: "https://user:secret@cdn.example.test/doc.pdf?X-Amz-Signature=abc123&token=leak#frag",
    });
    expect(line).toBe("Reading page: https://cdn.example.test/doc.pdf");
    expect(line).not.toContain("secret");
    expect(line).not.toContain("X-Amz-Signature");
    expect(line).not.toContain("token=");
    expect(line).not.toContain("abc123");
    expect(line).not.toContain("frag");
  });

  it("does not echo secrets from malformed web_fetch URLs", () => {
    const line = describeToolActivity("web_fetch", {
      url: "https://user:secret@[",
    });
    expect(line).toBe("Reading page: [invalid URL]");
    expect(line).not.toContain("secret");
    expect(line).not.toContain("user:");
  });

  it("falls back to the tool name", () => {
    expect(describeToolActivity("destination_write", undefined)).toBe("Using destination_write");
  });
});

describe("raceSignal", () => {
  it("resolves when the promise resolves before abort", async () => {
    const controller = new AbortController();
    const result = await raceSignal(Promise.resolve("ok"), controller.signal);
    expect(result).toBe("ok");
  });

  it("rejects when the promise rejects before abort", async () => {
    const controller = new AbortController();
    await expect(
      raceSignal(Promise.reject(new Error("tool failed")), controller.signal),
    ).rejects.toThrow("tool failed");
  });

  it("rejects with abort reason when aborted while promise is pending", async () => {
    const controller = new AbortController();
    let resolvePending: (val: string) => void;
    const pending = new Promise<string>((resolve) => {
      resolvePending = resolve;
    });

    const promise = raceSignal(pending, controller.signal);
    controller.abort(new Error("Interrupted by user speech"));

    await expect(promise).rejects.toThrow("Interrupted by user speech");
    resolvePending!("late");
  });

  it("rejects immediately if signal is already aborted", async () => {
    const controller = new AbortController();
    controller.abort(new Error("Already cancelled"));
    await expect(raceSignal(Promise.resolve("ok"), controller.signal)).rejects.toThrow(
      "Already cancelled",
    );
  });
});

describe("context window pruning and overflow recovery", () => {
  it("parses llama.cpp exceed_context_size_error and standard overflow errors", () => {
    const llamaCppError =
      '400: {"code":400,"message":"request (133699 tokens) exceeds the available context size (131072 tokens), try increasing it","type":"exceed_context_size_error","n_prompt_tokens":133699,"n_ctx":131072}';
    expect(parseContextOverflowDetails(llamaCppError)).toEqual({
      limitTokens: 131_072,
      requestedTokens: 133_699,
    });

    const plainLlamaError =
      "request (133699 tokens) exceeds the available context size (131072 tokens), try increasing it";
    expect(parseContextOverflowDetails(plainLlamaError)).toEqual({
      limitTokens: 131_072,
      requestedTokens: 133_699,
    });

    const anthropicError = "prompt is too long: 213462 tokens > 200000 maximum";
    expect(parseContextOverflowDetails(anthropicError)).toEqual({
      limitTokens: 200_000,
      requestedTokens: 213_462,
    });

    expect(parseContextOverflowDetails("429 Too many requests: rate limit exceeded")).toBeNull();
  });

  it("compacts older tool results and drops oldest history/tool rounds while keeping pairing intact", () => {
    const messages: AgentMessage[] = [
      {
        role: "system",
        content: "You are Trading-Bot.",
        toolsAdded: [{ name: "shell", description: "run shell", parameters: {} as never }],
        timestamp: 1,
      },
      {
        role: "user",
        content: `Old history 1: ${"h".repeat(9_000)}`,
        timestamp: 2,
      },
      {
        role: "user",
        content: `Old history 2: ${"h".repeat(9_000)}`,
        timestamp: 3,
      },
      {
        role: "user",
        content: "Fix the index error in backtest_p2.py.",
        timestamp: 4,
      },
      {
        role: "assistant",
        content: [
          { type: "thinking", thinking: "t".repeat(6_000) },
          {
            type: "toolCall",
            id: "call_1",
            name: "write_file",
            arguments: { path: "backtest_p2.py", content: "c".repeat(12_000) },
          },
        ],
        api: "openai-completions",
        provider: "local",
        model: "qwen",
        usage: {
          input: 0,
          output: 0,
          cacheRead: 0,
          cacheWrite: 0,
          totalTokens: 0,
          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
        },
        stopReason: "toolUse",
        timestamp: 5,
      },
      {
        role: "toolResult",
        toolCallId: "call_1",
        toolName: "write_file",
        content: [{ type: "text", text: "r".repeat(12_000) }],
        isError: false,
        timestamp: 6,
      },
      {
        role: "assistant",
        content: [
          {
            type: "toolCall",
            id: "call_2",
            name: "shell",
            arguments: { command: "python3 backtest_p2.py" },
          },
        ],
        api: "openai-completions",
        provider: "local",
        model: "qwen",
        usage: {
          input: 0,
          output: 0,
          cacheRead: 0,
          cacheWrite: 0,
          totalTokens: 0,
          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
        },
        stopReason: "toolUse",
        timestamp: 7,
      },
      {
        role: "toolResult",
        toolCallId: "call_2",
        toolName: "shell",
        content: [{ type: "text", text: "latest output: IndexError fixed" }],
        isError: false,
        timestamp: 8,
      },
    ];

    const pruned = pruneContextToWindow(messages, {
      contextWindow: 4_096,
      maxTokens: 1_024,
    });

    expect(estimateAgentContextTokens(pruned)).toBeLessThanOrEqual(
      maxInputTokensForWindow(4_096, 1_024),
    );
    expect(pruned[0]?.role).toBe("system");
    expect(
      pruned.some(
        (m) => m.role === "user" && m.content === "Fix the index error in backtest_p2.py.",
      ),
    ).toBe(true);
    const remainingToolCallIds = new Set(
      pruned
        .filter((m): m is Extract<AgentMessage, { role: "assistant" }> => m.role === "assistant")
        .flatMap((m) =>
          m.content
            .filter(
              (part): part is Extract<(typeof m.content)[number], { type: "toolCall" }> =>
                part.type === "toolCall",
            )
            .map((part) => part.id),
        ),
    );
    for (const msg of pruned) {
      if (msg.role === "toolResult") {
        expect(remainingToolCallIds.has(msg.toolCallId)).toBe(true);
      }
    }
  });

  it("retries transparently after a provider exceed_context_size_error", async () => {
    let effectiveWindow = 131_072;
    let calls = 0;

    const initialMessages: AgentMessage[] = [
      { role: "system", content: "You are a bot.", timestamp: 1 },
      { role: "user", content: `Prior history: ${"x".repeat(60_000)}`, timestamp: 2 },
      { role: "user", content: "Run the backtest.", timestamp: 3 },
    ];

    const stream = streamWithContextOverflowRecovery({
      models: {
        streamSimple: (_model, ctx) => {
          calls += 1;
          expect(ctx.messages.length).toBeGreaterThanOrEqual(2);
          const inner = createAssistantMessageEventStream();
          queueMicrotask(() => {
            if (calls === 1) {
              inner.push({
                type: "error",
                reason: "error",
                error: {
                  role: "assistant",
                  content: [],
                  api: "openai-completions",
                  provider: "local",
                  model: "qwen",
                  usage: {
                    input: 0,
                    output: 0,
                    cacheRead: 0,
                    cacheWrite: 0,
                    totalTokens: 0,
                    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
                  },
                  stopReason: "error",
                  errorMessage:
                    '400: {"code":400,"message":"request (133699 tokens) exceeds the available context size (131072 tokens), try increasing it","type":"exceed_context_size_error","n_prompt_tokens":133699,"n_ctx":131072}',
                  timestamp: Date.now(),
                },
              });
              inner.end();
              return;
            }
            const okMessage = {
              role: "assistant" as const,
              content: [{ type: "text" as const, text: "Backtest completed." }],
              api: "openai-completions" as const,
              provider: "local",
              model: "qwen",
              usage: {
                input: 1_000,
                output: 10,
                cacheRead: 0,
                cacheWrite: 0,
                totalTokens: 1_010,
                cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
              },
              stopReason: "stop" as const,
              timestamp: Date.now(),
            };
            inner.push({ type: "start", partial: okMessage });
            inner.push({
              type: "text_delta",
              contentIndex: 0,
              delta: "Backtest completed.",
              partial: okMessage,
            });
            inner.push({ type: "done", reason: "stop", message: okMessage });
            inner.end();
          });
          return inner;
        },
      },
      model: {
        id: "qwen",
        name: "qwen",
        api: "openai-completions",
        provider: "local",
        baseUrl: "http://127.0.0.1:11434/v1",
        reasoning: false,
        input: ["text"],
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
        contextWindow: 131_072,
        maxTokens: 4_096,
      },
      context: normalizeContext({ messages: initialMessages as never }),
      getEffectiveContextWindow: () => effectiveWindow,
      setEffectiveContextWindow: (next) => {
        effectiveWindow = next;
      },
    });

    const events = [];
    for await (const event of stream) {
      events.push(event.type);
    }
    const finalMsg = await stream.result();
    expect(calls).toBe(2);
    expect(effectiveWindow).toBeLessThan(131_072);
    expect(events).toEqual(["start", "text_delta", "done"]);
    expect(finalMsg.stopReason).toBe("stop");
  });
});