import { createHash } from "node:crypto";
import { and, eq, gt, inArray, ne } from "drizzle-orm";
import { environmentLeases, heartbeatRunEvents, heartbeatRuns, type Db } from "@paperclipai/db";
import { retiredReusableLeasePredecessors } from "../environment-lease-handoff.js";
import { hasRemoteTerminationReceipt } from "../remote-execution-termination.js";
import { hasConfirmedSandboxStopAndRetain } from "../sandbox-stop-and-retain.js";
import { LEGACY_WORKSPACE_RECOVERY_SCHEMA } from "../workspace-restore-recovery-state.js";
import { readNativeHarnessBackupStamp } from "./native-harness-backup-stamp.js";
import { nativeWorkspaceExportPendingForLease } from "./native-workspace-finalization-state.js";

/** Durable run events that prove a provider session did work in a run. */
export const NATIVE_PROVIDER_DURABLE_EVENT_TYPES = new Set([
  "harness.ready",
  "session.started",
  "session.resumed",
  "session.updated",
  "turn.started",
  "provider.event",
  "provider.rpc_result",
]);

export function nativeHarnessBackupManifestSha256(manifestBytes: Buffer | string): string {
  return `sha256:${createHash("sha256").update(manifestBytes).digest("hex")}`;
}

/**
 * A created sandbox normally refuses a per-turn backup that has no recorded
 * reusable lease, because the backup's source sandbox may still hold newer
 * provider state. That ambiguity is absent when Paperclip itself destroyed the
 * source sandbox: the provider confirmed the termination, the source lease's
 * stamp names exactly this backup, and no provider work for the session was
 * recorded after the backup completed. This is the same proof a confirmed
 * replacement relies on, reached through a lease that was retired first.
 *
 * A settled workspace-restore hold is the other retired source. The hold
 * changed the lease to `retain_on_failure`, which reuse never selects, and
 * stopped the sandbox under a stop-only receipt that permits no provider work.
 * Releases before the hold settlement restored warm reuse left such rows
 * stranded, so this also unsticks state saved by those releases.
 */
export async function retiredSourceAuthorizesNativeBackupRestore(db: Db, input: {
  companyId: string;
  currentRunId: string;
  normalizedSessionId: string;
  manifest: { sourceProviderLeaseId: string; runnerInstanceId: string; normalizedSessionId: string };
  manifestSha256: string;
}): Promise<boolean> {
  if (!input.manifest.sourceProviderLeaseId || input.manifest.normalizedSessionId !== input.normalizedSessionId) return false;
  const leases = await db.select().from(environmentLeases).where(and(
    eq(environmentLeases.companyId, input.companyId),
    eq(environmentLeases.providerLeaseId, input.manifest.sourceProviderLeaseId),
  ));
  const source = leases.filter(lease => lease.companyId === input.companyId
    && lease.providerLeaseId === input.manifest.sourceProviderLeaseId);
  const tips = source.filter(lease => !Object.hasOwn(lease.metadata ?? {}, "reusableLeaseSuccessor"));
  if (tips.length !== 1) return false;
  const lease = tips[0]!;
  if (!retiredReusableLeasePredecessors(lease, source.filter(row => row.id !== lease.id))) return false;
  const termination = lease.metadata?.remoteExecutionTermination as { state?: unknown } | undefined;
  const destroyed = ["expired", "failed"].includes(lease.status) && hasRemoteTerminationReceipt(lease)
    && termination?.state === "destroyed";
  const stamp = readNativeHarnessBackupStamp(lease.metadata?.nativeHarnessBackup, lease.providerLeaseId!);
  if (!stamp || stamp.manifestSha256 !== input.manifestSha256
    || stamp.normalizedSessionId !== input.normalizedSessionId
    || stamp.runnerInstanceId !== input.manifest.runnerInstanceId) return false;
  // A released lease can still be selected for resume; only a retired one is final.
  if (!destroyed && !await settledRestoreHoldRetiresSource(db, lease)) return false;
  // Provider work recorded after the checkpoint means the destroyed sandbox
  // held newer state than this backup. Restoring would silently rewind it.
  const later = await db.select({ runId: heartbeatRunEvents.runId }).from(heartbeatRunEvents)
    .innerJoin(heartbeatRuns, eq(heartbeatRuns.id, heartbeatRunEvents.runId))
    .where(and(
      eq(heartbeatRunEvents.companyId, input.companyId),
      eq(heartbeatRuns.companyId, input.companyId),
      eq(heartbeatRuns.nativeSessionId, input.normalizedSessionId),
      ne(heartbeatRuns.id, input.currentRunId),
      inArray(heartbeatRunEvents.eventType, [...NATIVE_PROVIDER_DURABLE_EVENT_TYPES]),
      gt(heartbeatRunEvents.createdAt, new Date(stamp.completedAt)),
    )).limit(1);
  return later.length === 0;
}

/**
 * A stopped, released lease is retired only when it is a workspace-restore hold
 * that nothing can select for resume: `retain_on_failure` with the hold's own
 * marker, a confirmed stop-and-retain receipt whose state is `stopped`, and no
 * unsaved writer or pending native workspace export bound to the allocation.
 * A stopped `reuse_by_environment` lease stays ambiguous and fails closed.
 */
async function settledRestoreHoldRetiresSource(db: Db, lease: typeof environmentLeases.$inferSelect): Promise<boolean> {
  const marker = lease.metadata?.workspaceRestoreRecovery as { schema?: unknown; runId?: unknown } | undefined;
  if (lease.leasePolicy !== "retain_on_failure" || lease.status !== "released" || !lease.heartbeatRunId
    || marker?.schema !== LEGACY_WORKSPACE_RECOVERY_SCHEMA || marker.runId !== lease.heartbeatRunId) return false;
  // Provider expiry cannot revive the allocation; it only makes the source
  // more final. Every other stop-and-retain receipt check stays strict.
  if (!hasConfirmedSandboxStopAndRetain({ ...lease, expiresAt: null })) return false;
  // Includes the dirty remote writer source (`preparedRemoteWorkspaceSource`).
  return !await nativeWorkspaceExportPendingForLease(db, lease);
}