import { createHash } from "node:crypto";
import {
  action,
  type ActionContext,
  type Action,
  type PendingApproval,
} from "@cloudflare/think";
import { callable } from "agents";
import * as Effect from "effect/Effect";
import type { FlexibleSchema, InferSchema } from "ai";
import {
  capabilityApprovalDecision,
  isCapabilityAction,
  type CapabilityApprovalRequest,
} from "../shared/capability-approval";
import type {
  CapabilityMetadata,
  ApprovalPolicy,
} from "../shared/capability-policy";
import { ActivityThink, type ToolActivityDescriptor } from "./tool-activity";

export interface CapabilityActionOptions<
  Schema extends FlexibleSchema,
  Output,
> {
  metadata: CapabilityMetadata;
  description?: string;
  activity?: Pick<ToolActivityDescriptor, "outcome" | "outputSummary"> & {
    version?: string;
  };
  inputSchema: Schema;
  policy: () => Promise<ApprovalPolicy | "unavailable">;
  execute: (
    input: InferSchema<Schema>,
    ctx: ActionContext,
    approvedOnce: boolean,
  ) => Promise<Output>;
}

const hash = (value: unknown) =>
  createHash("sha256").update(JSON.stringify(value)).digest("hex");
const approvalKey = (
  action: string,
  ctx: Pick<ActionContext, "requestId" | "toolCallId">,
  input: unknown,
) => hash([action, ctx.requestId, ctx.toolCallId, input]);

// Think owns the durable pending row, atomic claim, transcript replacement and
// continuation. The ephemeral grant binds this owner's decision to that exact
// native invocation and is consumed once, before guarded execution starts.
export class CapabilityApprovalThink<
  Env extends Cloudflare.Env,
> extends ActivityThink<Env> {
  private readonly approvalsInFlight = new Set<string>();
  private readonly approvedCalls = new Set<string>();
  private readonly capabilityMetadata = new WeakMap<
    Action,
    {
      metadata: CapabilityMetadata;
      activity?: CapabilityActionOptions<FlexibleSchema, unknown>["activity"];
    }
  >();
  async onStart() {
    super.onStart();
    for (const item of await this.pendingApprovals())
      if (
        item.source === "action" &&
        isCapabilityAction(item.descriptor.action)
      )
        this.restoreToolApproval(item.descriptor.toolCallId);
  }

  protected onActionApprovalDecision(event: {
    executionId: string;
    action: string;
    toolCallId: string;
    decision: "approve" | "deny";
  }) {
    super.onActionApprovalDecision(event);
    if (isCapabilityAction(event.action))
      this.recordToolApproval(event.toolCallId, event.decision);
  }

  protected resetTurnState() {
    this.approvedCalls.clear();
    super.resetTurnState();
  }

  protected getActionActivityDescriptor(action: Action) {
    const descriptor = super.getActionActivityDescriptor(action);
    const record = this.capabilityMetadata.get(action);
    if (!descriptor || !record) return descriptor;
    const { metadata, activity } = record;
    return {
      ...descriptor,
      ...(activity?.outputSummary
        ? { outputSummary: activity.outputSummary }
        : {}),
      outcome: (output: unknown) =>
        descriptor.outcome?.(output) === "pending"
          ? "pending"
          : (activity?.outcome?.(output) ??
            descriptor.outcome?.(output) ??
            "succeeded"),
      capability: {
        id: metadata.id,
        fingerprint: metadata.fingerprint,
        name: metadata.name,
        ...(activity?.version ? { version: activity.version } : {}),
        source: { ...metadata.source, name: metadata.sourceName },
      },
    };
  }

  protected capabilityAction<Schema extends FlexibleSchema, Output>(
    options: CapabilityActionOptions<Schema, Output>,
  ) {
    const name = `capability_${hash([options.metadata.source, options.metadata.id, options.metadata.fingerprint]).slice(0, 48)}`;
    const summary = `${options.metadata.name} from ${options.metadata.sourceName}: ${options.metadata.actionSummary}`;
    const capability = action({
      name,
      description: options.description ?? options.metadata.actionSummary,
      inputSchema: options.inputSchema,
      kind: "durable-pause",
      approvalSummary: summary,
      timeoutMs: 30_000,
      approval: async ({ input }) => {
        if (new TextEncoder().encode(JSON.stringify(input)).byteLength > 65_536)
          throw new Error(
            "Capability arguments exceed the approval review limit",
          );
        const policy = await options.policy();
        if (policy === "never" || policy === "unavailable")
          throw new Error("Capability is prohibited or unavailable");
        return policy === "ask";
      },
      execute: (input, ctx) =>
        options.execute(
          input,
          ctx,
          this.approvedCalls.delete(approvalKey(name, ctx, input)),
        ),
    });
    this.capabilityMetadata.set(capability, {
      metadata: options.metadata,
      activity: options.activity,
    });
    return capability;
  }

  @callable()
  async listCapabilityApprovals(): Promise<CapabilityApprovalRequest[]> {
    const pending = await this.pendingApprovals();
    return pending
      .filter(
        (item) =>
          item.source === "action" &&
          isCapabilityAction(item.descriptor.action),
      )
      .map((item) => {
        this.restoreToolApproval(item.descriptor.toolCallId);
        return {
          executionId: item.executionId,
          toolCallId: item.descriptor.toolCallId,
          action: item.descriptor.action,
          summary: item.descriptor.summary,
          arguments: item.descriptor.input,
        };
      });
  }

  @callable()
  decideCapabilityApproval(value: unknown) {
    return Effect.runPromise(
      Effect.gen({ self: this }, function* () {
        const decision = capabilityApprovalDecision.parse(value);
        if (this.approvalsInFlight.has(decision.executionId))
          return { resolved: false };
        this.approvalsInFlight.add(decision.executionId);
        let pending: PendingApproval | undefined;
        let key: string | undefined;
        try {
          pending = (yield* Effect.promise(() =>
            this.pendingApprovals(decision.executionId),
          )).find(
            (item) =>
              item.source === "action" &&
              isCapabilityAction(item.descriptor.action),
          );
          if (!pending) return { resolved: false };
          const descriptor = pending.descriptor;
          if (decision.decision === "approve") {
            key = approvalKey(descriptor.action, descriptor, descriptor.input);
            this.approvedCalls.add(key);
          }
          const result = yield* Effect.promise(() =>
            decision.decision === "approve"
              ? this.approveExecution(decision.executionId)
              : this.rejectExecution(
                  decision.executionId,
                  "Owner denied this action",
                ),
          );
          this.completeToolApproval(
            descriptor.toolCallId,
            decision.decision,
            result,
          );
          return { resolved: true };
        } finally {
          if (key) this.approvedCalls.delete(key);
          this.approvalsInFlight.delete(decision.executionId);
        }
      }),
    );
  }
}