import fs from "fs";
import path from "path";
import { buildRuntimeIds } from "./execution-config.mjs";
import { formatError } from "./formatting.mjs";
import {
  cleanupRuntimeInstanceContext,
  createRuntimeInstanceContext,
  ensureRuntimeInstanceReady,
} from "./runtime-contexts.mjs";

export function createRuntimeManager({ productDir, graphs, lifecycle, hooks = {}, runtimeOptions = {} }) {
  const graphByKey = new Map(graphs.map((graph) => [graph.key, graph]));
  const pools = new Map();
  const locks = new Map();
  const failedGraphs = new Map();
  let nextLeaseCounter = 1;
  const runtimeHooks = {
    cleanupRuntimeInstanceContext,
    createRuntimeInstanceContext,
    ensureRuntimeInstanceReady,
    sleep,
    options: runtimeOptions,
    ...hooks,
  };

  return {
    canAcquire(task) {
      if (failedGraphs.has(task.graphKey)) {
        return false;
      }
      const pool = getPool(pools, graphByKey, task, productDir, lifecycle);
      if (!locksAvailable(locks, task.locks || [])) {
        return false;
      }
      return claimableRuntimeSlot(pool, task) !== null;
    },
    async acquire(task) {
      const pool = getPool(pools, graphByKey, task, productDir, lifecycle);
      if (lifecycle.isStopRequested()) {
        throw lifecycle.signal.reason || new Error("testkit run interrupted");
      }

      const leaseId = `lease-${task.id}-${nextLeaseCounter}`;
      nextLeaseCounter += 1;

      if (!tryAcquireLocks(locks, task.locks || [], leaseId)) {
        throw new Error(`Task ${task.id} was claimed before its locks were available`);
      }

      const slot = claimRuntimeSlot(pool, task);
      if (!slot) {
        releaseLocks(locks, task.locks || [], leaseId);
        throw new Error(`Task ${task.id} was claimed before runtime capacity was available`);
      }

      try {
        const context = await getReadyContext(slot, productDir, task, lifecycle, runtimeHooks);
        const leaseDir = path.join(context.runtimeDir, "leases", leaseId);
        fs.mkdirSync(leaseDir, { recursive: true });
        return {
          leaseId,
          leaseDir,
          lockNames: task.locks || [],
          resourceCost: task.resourceCost || 1,
          slot,
          context,
        };
      } catch (error) {
        releaseRuntimeSlot(slot, task);
        releaseLocks(locks, task.locks || [], leaseId);
        cleanupLeaseDir({ leaseDir: path.join(slot.context?.runtimeDir || productDir, "leases", leaseId) });
        await drainRuntimeSlot(slot, lifecycle, runtimeHooks);
        if (!lifecycle.isStopRequested()) {
          markGraphFailed(failedGraphs, task.graphKey, formatError(error));
        }
        throw error;
      }
    },
    async release(lease, options = {}) {
      if (!lease?.slot) return;
      releaseLocks(locks, lease.lockNames || [], lease.leaseId);
      releaseRuntimeSlot(lease.slot, { resourceCost: lease.resourceCost || 1 });
      cleanupLeaseDir(lease);
      if (options.invalidate) {
        lease.slot.draining = true;
      }
      if (lease.slot.draining && lease.slot.activeLeaseCount === 0) {
        await drainRuntimeSlot(lease.slot, lifecycle, runtimeHooks);
      }
    },
    async cleanupAll() {
      for (const pool of pools.values()) {
        for (const slot of pool.slots) {
          slot.draining = true;
          await drainRuntimeSlot(slot, lifecycle, runtimeHooks);
        }
      }
    },
    getStats() {
      return [...pools.values()]
        .map((pool) => ({
          graphKey: pool.slots[0]?.graph.key || null,
          targetNames: [...(pool.slots[0]?.graph.targetNames || [])],
          maxConcurrentTasks: Number.isFinite(pool.slots[0]?.graph.maxConcurrentTasks)
            ? pool.slots[0]?.graph.maxConcurrentTasks
            : null,
          runtimeCount: pool.slots.length,
          runtimes: pool.slots.map((slot) => ({
            runtimeId: slot.runtimeId,
            peakLeaseCount: slot.peakLeaseCount,
            peakResourceUnits: slot.peakResourceUnits,
          })),
        }))
        .sort((left, right) => String(left.graphKey).localeCompare(String(right.graphKey)));
    },
    drainFailedTasks(queue) {
      const drained = [];
      for (let index = queue.length - 1; index >= 0; index -= 1) {
        const task = queue[index];
        const message = failedGraphs.get(task.graphKey);
        if (!message) continue;

        queue.splice(index, 1);
        drained.push({
          task,
          reason: `Not run because runtime graph failed to initialize: ${message}`,
          message,
          graph: graphByKey.get(task.graphKey) || {
            key: task.graphKey,
            targetNames: [task.targetName],
          },
        });
      }

      return drained.reverse();
    },
  };
}

function getPool(pools, graphByKey, task, productDir, lifecycle) {
  const existing = pools.get(task.graphKey);
  if (existing) return existing;

  const graph = graphByKey.get(task.graphKey);
  if (!graph) {
    throw new Error(`Unknown graph "${task.graphKey}"`);
  }

  const pool = {
    productDir,
    lifecycle,
    slots: buildRuntimeIds(graph.instanceCount).map((runtimeId) => ({
      graph,
      runtimeId,
      context: null,
      contextPromise: null,
      activeLeaseCount: 0,
      activeResourceUnits: 0,
      peakLeaseCount: 0,
      peakResourceUnits: 0,
      draining: false,
    })),
  };
  pools.set(task.graphKey, pool);
  return pool;
}

function claimRuntimeSlot(pool, task) {
  const resourceCost = task.resourceCost || 1;
  const available = pool.slots.filter((slot) => !slot.draining);
  if (available.length === 0) return null;

  const slot = [...available]
    .filter((candidate) => slotHasCapacity(candidate, resourceCost))
    .sort(
      (left, right) =>
        left.activeResourceUnits - right.activeResourceUnits ||
        left.activeLeaseCount - right.activeLeaseCount ||
        left.runtimeId.localeCompare(right.runtimeId)
    )[0];
  if (!slot) return null;
  slot.activeLeaseCount += 1;
  slot.activeResourceUnits += resourceCost;
  slot.peakLeaseCount = Math.max(slot.peakLeaseCount, slot.activeLeaseCount);
  slot.peakResourceUnits = Math.max(slot.peakResourceUnits, slot.activeResourceUnits);
  return slot;
}

function releaseRuntimeSlot(slot, task = {}) {
  slot.activeLeaseCount = Math.max(0, slot.activeLeaseCount - 1);
  slot.activeResourceUnits = Math.max(0, slot.activeResourceUnits - (task.resourceCost || 1));
}

async function getReadyContext(slot, productDir, task, lifecycle, runtimeHooks) {
  if (!slot.context) {
    slot.context = runtimeHooks.createRuntimeInstanceContext(slot.runtimeId, slot.graph, productDir);
    lifecycle.trackGraphContext(slot.context);
  }
  if (!slot.contextPromise) {
    slot.contextPromise = Promise.resolve(slot.context);
  }
  await slot.contextPromise;
  await runtimeHooks.ensureRuntimeInstanceReady(
    slot.context,
    task,
    lifecycle,
    runtimeHooks.options || {}
  );
  return slot.context;
}

async function drainRuntimeSlot(slot, lifecycle, runtimeHooks) {
  if (!slot.context || slot.activeLeaseCount > 0) return;
  await runtimeHooks.cleanupRuntimeInstanceContext(slot.context, lifecycle);
  slot.context = null;
  slot.contextPromise = null;
  slot.draining = false;
}

function tryAcquireLocks(lockMap, lockNames, leaseId) {
  const uniqueLocks = [...new Set(lockNames)].sort();
  for (const lockName of uniqueLocks) {
    if (lockMap.has(lockName)) {
      return false;
    }
  }
  for (const lockName of uniqueLocks) {
    lockMap.set(lockName, leaseId);
  }
  return true;
}

function locksAvailable(lockMap, lockNames) {
  return [...new Set(lockNames)].every((lockName) => !lockMap.has(lockName));
}

function claimableRuntimeSlot(pool, task) {
  const resourceCost = task.resourceCost || 1;
  return (
    pool.slots.find((slot) => !slot.draining && slotHasCapacity(slot, resourceCost)) || null
  );
}

function slotHasCapacity(slot, resourceCost) {
  return slot.activeResourceUnits + resourceCost <= slot.graph.maxConcurrentTasks;
}

function releaseLocks(lockMap, lockNames, leaseId) {
  for (const lockName of [...new Set(lockNames)].sort()) {
    if (lockMap.get(lockName) === leaseId) {
      lockMap.delete(lockName);
    }
  }
}

function markGraphFailed(failedGraphs, graphKey, message) {
  if (!graphKey || failedGraphs.has(graphKey)) return;
  failedGraphs.set(graphKey, message);
}

function sleep(ms) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

function cleanupLeaseDir(lease) {
  if (!lease?.leaseDir) return;
  fs.rmSync(lease.leaseDir, { recursive: true, force: true });
}
