import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { redactSecrets } from "../shared/secret-patterns.js";

function redactMessageContent(content: unknown): {
  content: unknown;
  found: string[];
} {
  if (typeof content === "string") {
    const { text, found } = redactSecrets(content);
    return { content: text, found };
  }

  if (Array.isArray(content)) {
    const found: string[] = [];
    const redacted = content.map((item) => {
      if (
        typeof item === "object" &&
        item !== null &&
        "type" in item &&
        item.type === "text" &&
        "text" in item &&
        typeof item.text === "string"
      ) {
        const result = redactSecrets(item.text);
        found.push(...result.found);
        return { ...item, text: result.text };
      }
      return item;
    });
    return { content: redacted, found };
  }

  return { content, found: [] };
}

type ContentBearingMessage = {
  role?: unknown;
  content?: unknown;
};

function isContentBearingMessage(
  message: unknown,
): message is ContentBearingMessage {
  return (
    typeof message === "object" &&
    message !== null &&
    "role" in message &&
    "content" in message
  );
}

function shouldRedactRole(role: unknown): role is string {
  return (
    role === "user" ||
    role === "assistant" ||
    role === "toolResult" ||
    role === "custom"
  );
}

export default function (pi: ExtensionAPI) {
  // 1. Intercept new user input before it hits the session
  pi.on("input", async (event, ctx) => {
    const { text, found } = redactSecrets(event.text);

    if (found.length > 0) {
      ctx.ui.notify(
        `🔒 Secret-Guard: Redacted ${found.length} secret(s) from your input`,
        "warning",
      );
      return { action: "transform", text };
    }

    return { action: "continue" };
  });

  // 2. Scrub secrets from context before sending to LLM
  pi.on("context", async (event) => {
    let modified = false;

    const messages = event.messages.map((message) => {
      if (
        !isContentBearingMessage(message) ||
        !shouldRedactRole(message.role)
      ) {
        return message;
      }

      const { content, found } = redactMessageContent(message.content);
      if (found.length === 0) {
        return message;
      }

      modified = true;
      return { ...message, content };
    }) as typeof event.messages;

    return modified ? { messages } : undefined;
  });

  // 3. Scrub secrets from tool call inputs (e.g. bash commands with env vars)
  pi.on("tool_call", async (event, ctx) => {
    if (event.toolName !== "bash") return;

    const command = event.input.command;
    if (typeof command !== "string") return;

    const { text, found } = redactSecrets(command);

    if (found.length > 0) {
      ctx.ui.notify(
        `🔒 Secret-Guard: Redacted secret from bash command`,
        "warning",
      );
      event.input.command = text;
    }
  });

  // 4. Scrub secrets from tool results
  pi.on("tool_result", async (event) => {
    const result = event.content;
    if (!Array.isArray(result)) return;

    let modified = false;
    const redacted = result.map((item) => {
      if (item.type === "text" && typeof item.text === "string") {
        const { text, found } = redactSecrets(item.text);
        if (found.length > 0) {
          modified = true;
          return { ...item, text };
        }
      }
      return item;
    });

    if (modified) {
      return { content: redacted };
    }
  });
}