import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { checkAction } from "../../src/core";
import { configLoader } from "../../src/shared/config";
import {
  createFeatureRequestPayload,
  emitActionBlocked,
  emitYoloChanged,
  GUARDRAILS_FEATURE_REGISTER_EVENT,
  GUARDRAILS_FEATURE_REQUEST_EVENT,
  type GuardrailsFeatureId,
  type GuardrailsFeatureRegisterPayload,
} from "../../src/shared/events";
import {
  GUARDRAILS_INHERITED_YOLO_ENTRY,
  GUARDRAILS_INHERITED_YOLO_ENV,
  GUARDRAILS_YOLO_ENTRY,
  inheritedYoloMode,
  initialYoloMode,
  restoreYoloMode,
} from "../../src/shared/yolo";
import { registerAuditLog } from "./audit";
import { registerGuardrailsExamplesCommand } from "./commands/examples";
import { registerGuardrailsOnboardingCommand } from "./commands/onboarding";
import { isOnboardingPending } from "./commands/onboarding/config";
import { registerGuardrailsSettings } from "./commands/settings";
import { registerYoloCommand, setYoloStatus } from "./commands/yolo";
import {
  BLOCKED_TOOLS,
  compilePolicies,
  createPolicyRules,
  protectionRank,
} from "./rules";
import { extractTargets } from "./targets";

export function setupPolicyHook(
  pi: ExtensionAPI,
  isYoloEnabled: () => boolean,
): void {
  pi.on("tool_call", async (event, ctx) => {
    if (isYoloEnabled()) return;
    const config = configLoader.getConfig();
    if (!config.enabled || !config.features.policies) return;

    const policies = compilePolicies(config.policies.rules)
      .filter((policy) => BLOCKED_TOOLS[policy.protection].has(event.toolName))
      .sort(
        (a, b) => protectionRank(b.protection) - protectionRank(a.protection),
      );
    if (policies.length === 0) return;

    const input = event.input as Record<string, unknown>;
    const targets = await extractTargets(
      { toolName: event.toolName, input },
      ctx.cwd,
      policies,
    );
    const rules = createPolicyRules(policies, ctx.cwd);

    for (const target of targets) {
      const safety = await checkAction(
        { kind: "file", path: target, origin: event.toolName },
        rules,
      );
      if (safety.kind === "safe") continue;

      emitActionBlocked(pi, {
        feature: "policies",
        action: safety.action,
        reason: safety.reason,
        block: { source: "policy", metadata: safety.metadata },
        context: { toolName: event.toolName, input },
      });
      return { block: true, reason: safety.reason };
    }
  });
}

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

  const loadedFeatures = new Set<GuardrailsFeatureId>(["policies"]);

  pi.events.on(GUARDRAILS_FEATURE_REGISTER_EVENT, (data: unknown) => {
    const payload = data as GuardrailsFeatureRegisterPayload;
    loadedFeatures.add(payload.feature.id);
  });

  registerGuardrailsSettings(pi, {
    getLoadedFeatures: () => loadedFeatures,
  });

  registerGuardrailsExamplesCommand(pi);
  if (isOnboardingPending(configLoader.getRawConfig("global"))) {
    registerGuardrailsOnboardingCommand(pi);
  }

  let yoloEnabled = false;
  const applyYolo = (enabled: boolean) => {
    yoloEnabled = enabled;
    process.env[GUARDRAILS_INHERITED_YOLO_ENV] = enabled ? "1" : "0";
    emitYoloChanged(pi, enabled);
  };

  registerYoloCommand(pi, {
    isEnabled: () => yoloEnabled,
    setEnabled: (enabled) => {
      pi.appendEntry(GUARDRAILS_YOLO_ENTRY, { enabled });
      applyYolo(enabled);
    },
  });
  setupPolicyHook(pi, () => yoloEnabled);
  registerAuditLog(pi);

  pi.on("session_start", (event, ctx) => {
    loadedFeatures.clear();
    loadedFeatures.add("policies");

    const restored = restoreYoloMode(ctx.sessionManager.getBranch());
    const isSubagent = process.env.PI_SUBAGENT_CHILD === "1";
    const inherited = isSubagent ? inheritedYoloMode() : undefined;
    const enabled = initialYoloMode({
      reason: event.reason,
      restored,
      inherited,
      isSubagent,
    });
    if (
      (event.reason === "new" || event.reason === "fork") &&
      inherited !== undefined
    ) {
      pi.appendEntry(GUARDRAILS_INHERITED_YOLO_ENTRY, {
        enabled: inherited,
        source: "subagent-environment",
      });
    }
    if (event.reason === "fork" && restored && inherited === undefined) {
      pi.appendEntry(GUARDRAILS_YOLO_ENTRY, { enabled: false });
    }
    applyYolo(enabled);
    setYoloStatus(ctx, enabled);

    pi.events.emit(
      GUARDRAILS_FEATURE_REQUEST_EVENT,
      createFeatureRequestPayload(),
    );
  });

  pi.on("session_tree", (_event, ctx) => {
    const enabled = restoreYoloMode(ctx.sessionManager.getBranch());
    applyYolo(enabled);
    setYoloStatus(ctx, enabled);
  });
}