import type { ExtensionAPI, ToolResultEvent } from "@mariozechner/pi-coding-agent";
import { mkdirSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { notify } from "../../src/shared/ui";
import { DEFAULT_CONFIG, loadOutputGuardConfig, resolveToolRule } from "./config";
import { measureContent, truncateContentBlocks } from "./truncate";
import type { OutputGuardConfig } from "./types";

export default function outputGuard(pi: ExtensionAPI) {
  let config: OutputGuardConfig = DEFAULT_CONFIG;

  pi.on("session_start", async (_event, ctx) => {
    const loaded = loadOutputGuardConfig(ctx.cwd);
    config = loaded.config;
  });

  pi.on("tool_result", (event: ToolResultEvent, ctx) => {
    if (!config.enabled) return;

    const toolName = event.toolName;

    // Skip exempt tools
    if (config.exempt.includes(toolName)) return;

    // Skip if no text content
    if (!event.content || !Array.isArray(event.content)) return;

    const textBlocks = event.content.filter(
      (b: any): b is { type: "text"; text: string } => b.type === "text" && typeof b.text === "string",
    );
    if (textBlocks.length === 0) return;

    const rule = resolveToolRule(config, toolName);
    const total = measureContent(textBlocks);
    if (total <= rule.maxChars) return;

    // Capture the full pre-truncation text BEFORE calling truncateContentBlocks
    // — that call mutates textBlocks in place, so reading them afterwards would
    // collect the already-truncated text.
    const fullText = textBlocks.map((b: any) => b.text).join("");

    // Save full content to a temp file before truncating, so the marker can
    // include the path and the file genuinely contains the original content.
    let savedPath: string | undefined;
    try {
      const timestamp = Date.now();
      const safeName = toolName.replace(/[^a-zA-Z0-9_-]/g, "_");
      const filename = `output-guard_${safeName}_${timestamp}.txt`;
      const dir = join(tmpdir(), "pi-output-guard");
      mkdirSync(dir, { recursive: true });
      savedPath = join(dir, filename);
      writeFileSync(savedPath, fullText, "utf8");
    } catch {
      savedPath = undefined;
    }

    // Truncate in place — event.content blocks are mutable. Pass savedPath so
    // it appears in the marker text the consumer (calling agent) sees, not
    // only in the notify() banner the operator sees.
    const info = truncateContentBlocks(textBlocks, rule, toolName, savedPath);
    if (!info) return;

    // Report truncation via return value so Pi picks up the modified content
    const newContent = event.content.map((block: any) => {
      if (block.type !== "text") return block;
      const match = textBlocks.find((tb: any) => tb === block);
      return match ? { ...block, text: match.text } : block;
    });

    if (config.notifications.showTruncation) {
      const pct = (info.removedPercent * 100).toFixed(0);
      let msg = `output-guard: ${toolName} output truncated ${info.originalChars.toLocaleString()} → ${info.keptChars.toLocaleString()} chars (${pct}% removed)`;
      if (savedPath) {
        msg += `\n  full output saved to: ${savedPath}`;
      }
      notify(ctx, msg, "warning");
    }

    return {
      content: newContent,
      details: {
        ...(typeof event.details === "object" && event.details !== null ? event.details : {}),
        outputGuard: {
          truncated: true,
          originalChars: info.originalChars,
          keptChars: info.keptChars,
          removedPercent: info.removedPercent,
          headRatio: info.headRatio,
          tool: info.tool,
          savedPath,
        },
      },
    };
  });
}