import { spawn } from "node:child_process";
import { createHash, randomUUID } from "node:crypto";
import { constants } from "node:fs";
import { lstat, mkdir, open, rename, rm } from "node:fs/promises";
import { isAbsolute, join, resolve } from "node:path";
import { Worker } from "node:worker_threads";

const IDENTITY_PATTERN = /^[^\p{Cc}]+$/u;
const FINGERPRINT_PATTERN = /^[a-f0-9]{64}$/u;
const DIRECTORY_OPEN_FLAGS = constants.O_RDONLY | constants.O_DIRECTORY | constants.O_NOFOLLOW;
const PROVISIONER_WORKER_SOURCE = `
const { parentPort, workerData } = require("node:worker_threads");
void (async () => {
  try {
    const module = await import(workerData.moduleUrl);
    const provision = module[workerData.exportName];
    if (typeof provision !== "function") throw new Error("missing provisioner export");
    await provision(workerData.input);
    parentPort.postMessage({ ok: true });
  } catch {
    parentPort.postMessage({ ok: false });
  }
})();
`;


export class AccountLeaseError extends Error {
  constructor(code, trace = null) {
    super(code);
    this.name = "AccountLeaseError";
    this.code = code;
    this.trace = trace;
  }

  toJSON() {
    return { name: this.name, code: this.code, trace: this.trace };
  }
}

const FLOCK_HOLDER_SOURCE = `
const parentPid = process.ppid;
process.stdout.write("READY\\n");
setInterval(() => {
  if (process.ppid !== parentPid) process.exit(1);
}, 250);
`;

async function startFlockAccountLock(lockPath, options, timeoutMs, timers) {
  const spawnProcess = options.spawnProcess ?? spawn;
  const child = spawnProcess(
    options.flockBinary ?? "flock",
    ["--no-fork", "--exclusive", "--nonblock", lockPath, process.execPath, "-e", FLOCK_HOLDER_SOURCE],
    { detached: true, stdio: ["ignore", "pipe", "ignore"] },
  );
  let exited = false;
  const exit = new Promise((resolveExit) => {
    child.once("exit", (code, signal) => {
      exited = true;
      resolveExit({ code, signal });
    });
  });
  const ready = new Promise((resolveReady, rejectReady) => {
    let output = "";
    child.once("error", () => rejectReady(new AccountLeaseError("account-lease-lock-failed")));
    child.stdout.on("data", (chunk) => {
      output += String(chunk);
      if (output.includes("READY\n")) resolveReady();
    });
    exit.then(({ code }) => {
      if (!output.includes("READY\n")) {
        const reason = code === 1 ? "account-lease-collision" : "account-lease-lock-failed";
        rejectReady(new AccountLeaseError(reason));
      }
    });
  });
  let startupTimer;
  const startupTimeout = new Promise((_, reject) => {
    startupTimer = timers.setTimeout(
      () => reject(new AccountLeaseError("account-lease-lock-timeout")),
      timeoutMs,
    );
  });
  try {
    await Promise.race([ready, startupTimeout]);
  } catch (error) {
    if (!exited) {
      try {
        process.kill(-child.pid, "SIGKILL");
      } catch (killError) {
        if (killError?.code !== "ESRCH") throw new AccountLeaseError("account-lease-lock-failed");
      }
    }
    throw error;
  } finally {
    timers.clearTimeout(startupTimer);
  }
  child.unref();
  child.stdout.unref();
  return {
    assertAlive() {
      if (exited) throw new AccountLeaseError("account-lease-lock-lost");
    },
    async release() {
      if (exited) return;
      try {
        process.kill(-child.pid, "SIGKILL");
      } catch (error) {
        if (error?.code !== "ESRCH") throw new AccountLeaseError("account-lease-lock-release-failed");
      }
      await exit;
    },
  };
}

function requireText(value, code, maxLength = 256) {
  if (typeof value !== "string" || value.length === 0 || value.length > maxLength || !IDENTITY_PATTERN.test(value)) {
    throw new AccountLeaseError(code);
  }
  return value;
}

function validateResolvedAccount(value) {
  if (!value || typeof value !== "object" || Array.isArray(value)) {
    throw new AccountLeaseError("account-ref-unresolved");
  }
  const expectedIdentityValue = value.expectedIdentity;
  const accountFingerprintValue = value.accountFingerprint;
  const accountHandleValue = value.accountHandle;
  const concurrentSafeValue = value.concurrentSafe;
  const expectedIdentity = requireText(expectedIdentityValue, "account-expected-identity-invalid");
  if (typeof accountFingerprintValue !== "string" || !FINGERPRINT_PATTERN.test(accountFingerprintValue)) {
    throw new AccountLeaseError("account-fingerprint-invalid");
  }
  return {
    expectedIdentity,
    accountFingerprint: accountFingerprintValue,
    accountHandle: accountHandleValue,
    concurrentSafe: concurrentSafeValue === true,
  };
}

function observedIdentities(result) {
  if (typeof result === "string") return [result];
  if (!result || typeof result !== "object" || Array.isArray(result)) return [];
  const hasIdentity = Object.hasOwn(result, "identity");
  const hasIdentities = Object.hasOwn(result, "identities");
  if (hasIdentity && hasIdentities) return null;
  if (hasIdentity) {
    const identity = result.identity;
    return [identity];
  }
  if (hasIdentities) {
    const identities = result.identities;
    return Array.isArray(identities) ? Array.from(identities) : [];
  }
  return [];
}

function validateObservedIdentity(result) {
  const identities = observedIdentities(result);
  if (identities === null || identities.length > 1) throw new AccountLeaseError("account-identity-ambiguous");
  if (identities.length === 0) throw new AccountLeaseError("account-identity-unavailable");
  try {
    return requireText(identities[0], "account-identity-unavailable");
  } catch {
    throw new AccountLeaseError("account-identity-unavailable");
  }
}

function digest(value) {
  return createHash("sha256").update(value).digest("hex");
}

function descriptorIdentity(stat) {
  return { dev: stat.dev, ino: stat.ino };
}

function sameIdentity(left, right) {
  return left.dev === right.dev && left.ino === right.ino;
}

function descriptorPath(handle, name = "") {
  return name === "" ? `/proc/self/fd/${handle.fd}` : `/proc/self/fd/${handle.fd}/${name}`;
}

async function closeHandle(handle) {
  if (handle !== null) await handle.close();
}

async function closeHandles(handles) {
  const results = await Promise.allSettled(handles.map(closeHandle));
  if (results.some((result) => result.status === "rejected")) {
    throw new AccountLeaseError("account-lease-close-failed");
  }
}

async function readIdentity(fs, path, lostCode) {
  try {
    return descriptorIdentity(await fs.lstat(path));
  } catch {
    throw new AccountLeaseError(lostCode);
  }
}

async function assertPathIdentity(fs, path, expected, lostCode) {
  const actual = await readIdentity(fs, path, lostCode);
  if (!sameIdentity(actual, expected)) throw new AccountLeaseError(lostCode);
}

async function quarantinePath(fs, path, lostCode) {
  const quarantinedPath = `${path}.quarantine-${randomUUID()}`;
  try {
    await fs.rename(path, quarantinedPath);
  } catch {
    throw new AccountLeaseError(lostCode);
  }
  return quarantinedPath;
}

async function restoreQuarantine(fs, quarantinedPath, path) {
  try {
    await fs.rename(quarantinedPath, path);
  } catch {
    throw new AccountLeaseError("account-lease-cleanup-failed");
  }
}

async function removeOwnedPath(fs, path, expected, lostCode, cleanupState = null) {
  let quarantinedPath = cleanupState?.quarantinedPath ?? null;
  if (quarantinedPath === null) {
    quarantinedPath = await quarantinePath(fs, path, lostCode);
    try {
      await assertPathIdentity(fs, quarantinedPath, expected, lostCode);
    } catch (error) {
      await restoreQuarantine(fs, quarantinedPath, path);
      throw error;
    }
    if (cleanupState !== null) cleanupState.quarantinedPath = quarantinedPath;
  }
  await fs.rm(quarantinedPath, { recursive: true, force: true });
  if (cleanupState !== null) cleanupState.quarantinedPath = null;
}

function startProvisionerWorker(moduleUrl, exportName, input, WorkerClass) {
  const worker = new WorkerClass(PROVISIONER_WORKER_SOURCE, {
    eval: true,
    stderr: true,
    stdout: true,
    workerData: { moduleUrl, exportName, input },
  });
  worker.stdout?.resume();
  worker.stderr?.resume();
  let report = null;
  const completion = new Promise((resolveCompletion, rejectCompletion) => {
    worker.on("message", (message) => {
      report = message?.ok === true ? "ok" : "failed";
    });
    worker.once("error", () => {
      report = "failed";
    });
    worker.once("exit", (code) => {
      if (code === 0 && report === "ok") resolveCompletion();
      else rejectCompletion(new AccountLeaseError("account-provision-failed"));
    });
  });
  return {
    completion,
    async terminate(timeoutMs) {
      let terminationError = null;
      for (let attempt = 0; attempt < 2 && worker.threadId !== -1; attempt += 1) {
        let timer;
        try {
          await Promise.race([
            worker.terminate(),
            new Promise((_, reject) => {
              timer = setTimeout(() => reject(new AccountLeaseError("account-provision-termination-timeout")), timeoutMs);
            }),
          ]);
        } catch (error) {
          terminationError ??= error;
        } finally {
          if (timer !== undefined) clearTimeout(timer);
        }
      }
      if (worker.threadId !== -1) {
        throw new AccountLeaseError("account-provision-death-unconfirmed");
      }
      if (terminationError !== null) throw terminationError;
    },
  };
}

async function runBounded(operation, timeoutMs, timeoutCode, timers) {
  const controller = new AbortController();
  let timer;
  const timeout = new Promise((_, reject) => {
    timer = timers.setTimeout(() => {
      controller.abort();
      reject(new AccountLeaseError(timeoutCode));
    }, timeoutMs);
  });
  try {
    return await Promise.race([Promise.resolve().then(() => operation(controller.signal)), timeout]);
  } finally {
    timers.clearTimeout(timer);
  }
}

export function createAccountLeaseManager(options) {
  if (!options || typeof options !== "object") throw new AccountLeaseError("account-lease-options-invalid");
  const leaseRoot = options.leaseRoot;
  if (typeof leaseRoot !== "string" || !isAbsolute(leaseRoot) || resolve(leaseRoot) !== leaseRoot) {
    throw new AccountLeaseError("account-lease-root-invalid");
  }
  if (typeof options.resolveAccount !== "function") throw new AccountLeaseError("account-resolver-required");
  if (typeof options.snapshotProcessIdentity !== "function") throw new AccountLeaseError("account-process-snapshot-required");
  if (typeof options.verifyIdentity !== "function") throw new AccountLeaseError("account-identity-verifier-required");
  const provisionerModuleUrl = requireText(options.provisionerModuleUrl, "account-provisioner-module-required", 4096);
  try {
    new URL(provisionerModuleUrl);
  } catch {
    throw new AccountLeaseError("account-provisioner-module-invalid");
  }
  const provisionerExport = options.provisionerExport ?? "provisionAccount";
  requireText(provisionerExport, "account-provisioner-export-invalid", 128);
  const WorkerClass = options.Worker ?? Worker;
  if (typeof WorkerClass !== "function") throw new AccountLeaseError("account-provisioner-worker-invalid");

  const operationTimeoutMs = options.operationTimeoutMs ?? 30_000;
  const provisionTimeoutMs = options.provisionTimeoutMs ?? operationTimeoutMs;
  const identityTimeoutMs = options.identityTimeoutMs ?? operationTimeoutMs;
  const lockTimeoutMs = options.lockTimeoutMs ?? 5_000;
  for (const timeout of [operationTimeoutMs, provisionTimeoutMs, identityTimeoutMs, lockTimeoutMs]) {
    if (!Number.isInteger(timeout) || timeout < 1 || timeout > 300_000) {
      throw new AccountLeaseError("account-operation-timeout-invalid");
    }
  }
  const fs = options.fs ?? { lstat, mkdir, open, rename, rm };
  if ([fs.lstat, fs.mkdir, fs.open, fs.rename, fs.rm].some((fn) => typeof fn !== "function")) {
    throw new AccountLeaseError("account-lease-filesystem-invalid");
  }
  const timers = {
    setTimeout: options.setTimeout ?? setTimeout,
    clearTimeout: options.clearTimeout ?? clearTimeout,
  };
  if (typeof timers.setTimeout !== "function" || typeof timers.clearTimeout !== "function") {
    throw new AccountLeaseError("account-lease-timers-invalid");
  }
  const acquireAccountLock = options.acquireAccountLock
    ?? ((lockPath) => startFlockAccountLock(lockPath, options, lockTimeoutMs, timers));
  if (typeof acquireAccountLock !== "function") throw new AccountLeaseError("account-lock-provider-invalid");

  const runsRoot = join(leaseRoot, "runs");
  const locksRoot = join(leaseRoot, "locks");

  return Object.freeze({
    async acquire(input) {
      const runId = requireText(input?.runId, "account-lease-run-id-invalid", 128);
      const accountRef = requireText(input?.accountRef, "account-ref-invalid");

      let resolvedAccount;
      try {
        resolvedAccount = await runBounded(
          (signal) => options.resolveAccount(accountRef, { signal }),
          operationTimeoutMs,
          "account-ref-unavailable",
          timers,
        );
      } catch {
        throw new AccountLeaseError("account-ref-unavailable");
      }
      let account;
      try {
        account = validateResolvedAccount(resolvedAccount);
      } catch (error) {
        const allowedCodes = new Set([
          "account-ref-unresolved",
          "account-expected-identity-invalid",
          "account-fingerprint-invalid",
        ]);
        const code = error instanceof AccountLeaseError && allowedCodes.has(error.code)
          ? error.code
          : "account-ref-unavailable";
        throw new AccountLeaseError(code);
      }
      const publicAccountRef = `sha256:${digest(`account-ref\0${accountRef}`)}`;
      const runName = digest(`account-run\0${runId}\0${accountRef}`);
      const lockName = account.concurrentSafe ? null : digest(`account-lock\0${account.accountFingerprint}`);
      const codingAgentDir = join(runsRoot, runName);

      await fs.mkdir(leaseRoot, { recursive: true, mode: 0o700 });
      await fs.mkdir(runsRoot, { recursive: true, mode: 0o700 });
      await fs.mkdir(locksRoot, { recursive: true, mode: 0o700 });

      let runsHandle = null;
      let locksHandle = null;
      let runHandle = null;
      let accountLock = null;
      let runIdentity = null;
      let locksIdentity = null;
      try {
        runsHandle = await fs.open(runsRoot, DIRECTORY_OPEN_FLAGS);
        locksHandle = await fs.open(locksRoot, DIRECTORY_OPEN_FLAGS);
        locksIdentity = descriptorIdentity(await locksHandle.stat());
      } catch {
        await closeHandles([runHandle, runsHandle, locksHandle]);
        throw new AccountLeaseError("account-lease-root-unavailable");
      }

      const runPath = descriptorPath(runsHandle, runName);
      const lockPath = lockName === null
        ? null
        : `/proc/${process.pid}/fd/${locksHandle.fd}/${lockName}.lock`;

      if (lockPath !== null) {
        try {
          await assertPathIdentity(fs, locksRoot, locksIdentity, "account-lease-lock-root-lost");
          accountLock = await acquireAccountLock(lockPath);
          await assertPathIdentity(fs, locksRoot, locksIdentity, "account-lease-lock-root-lost");
          accountLock.assertAlive?.();
        } catch (error) {
          if (accountLock !== null) await accountLock.release().catch(() => undefined);
          await closeHandles([runsHandle, locksHandle]);
          if (error instanceof AccountLeaseError) throw new AccountLeaseError(error.code);
          throw new AccountLeaseError("account-lease-lock-failed");
        }
      }

      try {
        await fs.mkdir(runPath, { mode: 0o700 });
        runHandle = await fs.open(runPath, DIRECTORY_OPEN_FLAGS);
        runIdentity = descriptorIdentity(await runHandle.stat());
        await assertPathIdentity(fs, codingAgentDir, runIdentity, "account-lease-directory-lost");
      } catch (error) {
        if (accountLock !== null) await accountLock.release();
        await closeHandles([runHandle, runsHandle, locksHandle]);
        if (error?.code === "EEXIST") throw new AccountLeaseError("account-lease-collision");
        if (error instanceof AccountLeaseError) throw error;
        throw new AccountLeaseError("account-lease-directory-failed");
      }

      let provisioner = null;
      let provisionTimer = null;
      const provisionTimeout = new AccountLeaseError("account-provision-timeout");
      try {
        provisioner = startProvisionerWorker(
          provisionerModuleUrl,
          provisionerExport,
          {
            runId,
            accountRef,
            codingAgentDir: descriptorPath(runHandle),
            canonicalCodingAgentDir: codingAgentDir,
            provisioningDir: descriptorPath(runHandle),
            accountHandle: account.accountHandle,
          },
          WorkerClass,
        );
        const timeout = new Promise((_, reject) => {
          provisionTimer = timers.setTimeout(() => reject(provisionTimeout), provisionTimeoutMs);
        });
        await Promise.race([provisioner.completion, timeout]);
        await assertPathIdentity(fs, codingAgentDir, runIdentity, "account-lease-directory-lost");
        accountLock?.assertAlive?.();
      } catch (error) {
        let terminationFailed = false;
        let deathUnconfirmed = false;
        if (error === provisionTimeout) {
          try {
            await provisioner.terminate(Math.min(provisionTimeoutMs, 5_000));
          } catch (terminationError) {
            if (terminationError instanceof AccountLeaseError
              && terminationError.code === "account-provision-death-unconfirmed") {
              deathUnconfirmed = true;
            } else {
              terminationFailed = true;
            }
          }
        }
        if (deathUnconfirmed) {
          const cleanupAfterDeath = async () => {
            try {
              await removeOwnedPath(fs, runPath, runIdentity, "account-lease-directory-lost");
              if (accountLock !== null) await accountLock.release();
            } finally {
              await closeHandles([runHandle, runsHandle, locksHandle]);
            }
          };
          void provisioner.completion.then(cleanupAfterDeath, cleanupAfterDeath).catch(() => undefined);
          throw new AccountLeaseError("account-provision-capability-quarantined");
        }
        if (error instanceof AccountLeaseError && error.code === "account-lease-directory-lost") {
          if (accountLock !== null) await accountLock.release().catch(() => undefined);
          await closeHandles([runHandle, runsHandle, locksHandle]);
          throw new AccountLeaseError("account-lease-directory-lost");
        }
        let cleanupFailed = false;
        try {
          await removeOwnedPath(fs, runPath, runIdentity, "account-lease-directory-lost");
          if (accountLock !== null) await accountLock.release();
        } catch {
          cleanupFailed = true;
        } finally {
          await closeHandles([runHandle, runsHandle, locksHandle]);
        }
        if (cleanupFailed) throw new AccountLeaseError("account-lease-cleanup-failed");
        if (terminationFailed) throw new AccountLeaseError("account-provision-termination-failed");
        if (error instanceof AccountLeaseError && error.code === "account-lease-lock-lost") throw error;
        if (error === provisionTimeout) throw provisionTimeout;
        throw new AccountLeaseError("account-provision-failed");
      } finally {
        if (provisionTimer !== null) timers.clearTimeout(provisionTimer);
      }

      let observedIdentity = null;
      let state = "active";
      let releasePromise = null;
      let verificationPromise = null;
      let verificationProcessIdentity = null;
      let runRemoved = false;
      const runCleanupState = { quarantinedPath: null };
      let lockReleased = false;

      const trace = () => Object.freeze({
        schemaVersion: 1,
        runId,
        accountRef: publicAccountRef,
        codingAgentDir,
        requestedIdentity: account.expectedIdentity,
        expectedIdentity: account.expectedIdentity,
        observedIdentity,
        accountFingerprint: account.accountFingerprint,
      });

      const lease = {
        schemaVersion: 1,
        runId,
        accountRef: publicAccountRef,
        codingAgentDir,
        expectedIdentity: account.expectedIdentity,
        accountFingerprint: account.accountFingerprint,
        env: Object.freeze({ PI_CODING_AGENT_DIR: codingAgentDir }),
        get observedIdentity() {
          return observedIdentity;
        },
        trace,
        async verifyAfterStartup(processHandle) {
          if (state !== "active") {
            const code = state === "released" ? "account-lease-released" : "account-lease-releasing";
            throw new AccountLeaseError(code, trace());
          }
          if (processHandle === null || processHandle === undefined) {
            throw new AccountLeaseError("account-process-unavailable", trace());
          }
          let processIdentity;
          try {
            processIdentity = requireText(
              options.snapshotProcessIdentity(processHandle),
              "account-process-identity-invalid",
              512,
            );
          } catch {
            throw new AccountLeaseError("account-process-identity-invalid", trace());
          }
          try {
            accountLock?.assertAlive?.();
          } catch {
            throw new AccountLeaseError("account-lease-lock-lost", trace());
          }
          if (verificationPromise !== null && processIdentity !== verificationProcessIdentity) {
            throw new AccountLeaseError("account-process-mismatch", trace());
          }
          if (verificationPromise === null) {
            verificationProcessIdentity = processIdentity;
            verificationPromise = (async () => {
              let result;
              try {
                result = await runBounded(
                  (abortSignal) => options.verifyIdentity({
                    processIdentity,
                    runId,
                    accountRef,
                    codingAgentDir,
                    expectedIdentity: account.expectedIdentity,
                    abortSignal,
                  }),
                  identityTimeoutMs,
                  "account-identity-unavailable",
                  timers,
                );
              } catch {
                throw new AccountLeaseError("account-identity-unavailable", trace());
              }
              try {
                observedIdentity = validateObservedIdentity(result);
              } catch (error) {
                const code = error instanceof AccountLeaseError ? error.code : "account-identity-unavailable";
                throw new AccountLeaseError(code, trace());
              }
              if (observedIdentity !== account.expectedIdentity) {
                throw new AccountLeaseError("account-identity-mismatch", trace());
              }
              return trace();
            })();
          }
          return verificationPromise;
        },
        release() {
          if (state === "released") return Promise.resolve();
          if (state === "releasing") return releasePromise;
          state = "releasing";
          releasePromise = (async () => {
            if (verificationPromise !== null) await verificationPromise.catch(() => undefined);
            if (!runRemoved) {
              await removeOwnedPath(
                fs,
                runPath,
                runIdentity,
                "account-lease-directory-lost",
                runCleanupState,
              );
              runRemoved = true;
            }
            if (accountLock !== null && !lockReleased) {
              await accountLock.release();
              lockReleased = true;
            }
            await closeHandles([runHandle, runsHandle, locksHandle]);
            state = "released";
          })().catch((error) => {
            state = "release-failed";
            releasePromise = null;
            if (error instanceof AccountLeaseError) throw error;
            throw new AccountLeaseError("account-lease-release-failed", trace());
          });
          return releasePromise;
        },
        toJSON() {
          return trace();
        },
      };

      return Object.freeze(lease);
    },
  });
}
