import { randomUUID } from "node:crypto";
import {
  type ExtensionAPI,
  isToolCallEventType,
} from "@earendil-works/pi-coding-agent";
import { checkAction } from "../../src/core";
import { configLoader } from "../../src/shared/config";
import {
  createFeatureRegisterPayload,
  emitActionBlocked,
  emitActionPrompted,
  emitRiskDetected,
  GUARDRAILS_APPROVAL_BROKER_REGISTER_EVENT,
  GUARDRAILS_APPROVAL_REPLY_EVENT,
  GUARDRAILS_APPROVAL_REQUEST_EVENT,
  GUARDRAILS_APPROVAL_SESSION_GRANT_EVENT,
  GUARDRAILS_FEATURE_REGISTER_EVENT,
  GUARDRAILS_FEATURE_REQUEST_EVENT,
  GUARDRAILS_YOLO_CHANGED_EVENT,
  type GuardrailsApprovalBrokerRegisterPayload,
  type GuardrailsApprovalDecision,
  type GuardrailsApprovalReplyPayload,
  type GuardrailsApprovalSessionGrantPayload,
  type GuardrailsYoloChangedPayload,
} from "../../src/shared/events";
import { saveCommandSessionGrant } from "./grants";
import { createPermissionGateConfirmComponent } from "./prompt";
import {
  createPermissionGateRule,
  formatAutoDenyReason,
  matchCommandPattern,
} from "./rules";

export default async function permissionGate(pi: ExtensionAPI) {
  await configLoader.load();

  let yoloEnabled = false;
  let approvalBrokerRegistered = false;
  const sessionGrants = new Set<string>();
  pi.events.on(GUARDRAILS_APPROVAL_SESSION_GRANT_EVENT, (data: unknown) => {
    const payload = data as GuardrailsApprovalSessionGrantPayload | undefined;
    if (payload?.version === 1 && payload.command)
      sessionGrants.add(payload.command);
  });
  pi.events.on(GUARDRAILS_APPROVAL_BROKER_REGISTER_EVENT, (data: unknown) => {
    const payload = data as GuardrailsApprovalBrokerRegisterPayload | undefined;
    if (payload?.version === 1 && payload.brokerId)
      approvalBrokerRegistered = true;
  });
  pi.events.on(GUARDRAILS_YOLO_CHANGED_EVENT, (data: unknown) => {
    const payload = data as GuardrailsYoloChangedPayload | undefined;
    if (typeof payload?.enabled === "boolean") yoloEnabled = payload.enabled;
  });

  pi.events.on(GUARDRAILS_FEATURE_REQUEST_EVENT, () => {
    pi.events.emit(
      GUARDRAILS_FEATURE_REGISTER_EVENT,
      createFeatureRegisterPayload("permissionGate"),
    );
  });

  pi.on("tool_call", async (event, ctx) => {
    if (yoloEnabled) return;
    const config = configLoader.getConfig();
    if (!config.enabled || !config.features.permissionGate) return;
    if (!isToolCallEventType("bash", event)) return;

    const command = event.input.command;
    const action = { kind: "command" as const, command, origin: "bash" };
    if (
      sessionGrants.has(command) ||
      matchCommandPattern(command, config.permissionGate.allowedPatterns)
    ) {
      return;
    }

    const autoDenyMatch = matchCommandPattern(
      command,
      config.permissionGate.autoDenyPatterns,
    );

    if (autoDenyMatch) {
      const reason = formatAutoDenyReason(autoDenyMatch);

      emitActionBlocked(pi, {
        feature: "permissionGate",
        action,
        reason,
        block: { source: "permission", metadata: autoDenyMatch },
        context: { toolName: "bash", input: event.input },
      });

      return { block: true, reason };
    }

    const safety = await checkAction(action, [
      createPermissionGateRule({
        patterns: config.permissionGate.patterns,
        useBuiltinMatchers: config.permissionGate.useBuiltinMatchers,
      }),
    ]);
    if (safety.kind === "safe") return;

    emitRiskDetected(pi, {
      feature: "permissionGate",
      risk: safety,
      context: { toolName: "bash", input: event.input },
    });

    if (!config.permissionGate.requireConfirmation) {
      ctx.ui.notify(`Dangerous command detected: ${safety.reason}`, "warning");
      return;
    }

    let result: GuardrailsApprovalDecision | undefined;
    let brokerFailureReason: string | undefined;
    if (!ctx.hasUI) {
      if (approvalBrokerRegistered) {
        const requestId = randomUUID();
        const timeoutMs =
          Number(process.env.PI_GUARDRAILS_APPROVAL_TIMEOUT_MS) ||
          10 * 60 * 1000;
        result = await new Promise<GuardrailsApprovalDecision | undefined>(
          (resolve) => {
            const timer = setTimeout(() => resolve(undefined), timeoutMs);
            const unsubscribe = pi.events.on(
              GUARDRAILS_APPROVAL_REPLY_EVENT,
              (data: unknown) => {
                const reply = data as
                  | GuardrailsApprovalReplyPayload
                  | undefined;
                if (reply?.version !== 1 || reply.requestId !== requestId)
                  return;
                clearTimeout(timer);
                unsubscribe?.();
                brokerFailureReason = reply.reason;
                resolve(reply.decision);
              },
            );
            pi.events.emit(GUARDRAILS_APPROVAL_REQUEST_EVENT, {
              version: 1,
              requestId,
              command,
              reason: safety.reason,
              metadata: safety.metadata,
            });
          },
        );
      }
      if (!result) {
        const reason = approvalBrokerRegistered
          ? `Dangerous command blocked (${brokerFailureReason ?? "approval timed out"}): ${safety.reason}`
          : `Dangerous command blocked (no UI to confirm): ${safety.reason}`;
        emitActionBlocked(pi, {
          feature: "permissionGate",
          action: safety.action,
          reason,
          block: { source: "nonInteractive", metadata: safety.metadata },
          context: { toolName: "bash", input: event.input },
        });
        return { block: true, reason };
      }
    }

    emitActionPrompted(pi, {
      feature: "permissionGate",
      action: safety.action,
      reason: safety.reason,
      prompt: {
        kind: "permission",
        metadata: safety.metadata,
      },
      context: { toolName: "bash", input: event.input },
    });

    if (ctx.hasUI) {
      result = await ctx.ui.custom<GuardrailsApprovalDecision>(
        createPermissionGateConfirmComponent(command, safety.reason),
      );
    }

    if (ctx.hasUI && result === undefined) {
      const selection = await ctx.ui.select(
        `Dangerous command: ${safety.reason}`,
        ["Allow once", "Allow for session", "Deny", "Decline and stop"],
      );
      if (selection === "Allow once") result = "allow";
      else if (selection === "Allow for session") result = "allow-session";
      else if (selection === "Decline and stop") result = "stop";
      else result = "deny";
    }

    if (result === "allow") return;
    if (result === "allow-session") {
      sessionGrants.add(command);
      pi.events.emit(GUARDRAILS_APPROVAL_SESSION_GRANT_EVENT, {
        version: 1,
        command,
      } satisfies GuardrailsApprovalSessionGrantPayload);
      await saveCommandSessionGrant(command);
      return;
    }

    if (result === "stop") {
      const reason = "User declined and stopped dangerous command";
      emitActionBlocked(pi, {
        feature: "permissionGate",
        action: safety.action,
        reason,
        block: { source: "user-stop", metadata: safety.metadata },
        context: { toolName: "bash", input: event.input },
      });
      ctx.abort();
      return { block: true, reason };
    }

    const reason = brokerFailureReason ?? "User denied dangerous command";
    emitActionBlocked(pi, {
      feature: "permissionGate",
      action: safety.action,
      reason,
      block: { source: "user", metadata: safety.metadata },
      context: { toolName: "bash", input: event.input },
    });
    return { block: true, reason };
  });
}