/**
 * Cleanup and shutdown handling for workflow runs.
 *
 * - Kills the sandbox child process on session_shutdown
 * - Adapted for Windows process tree termination via taskkill
 * - Background follow-up sends a single followUp message via pi.sendUserMessage
 * - Cleans up stale run artifacts
 */

import type { WorkflowProgress } from "./types.ts";
import { killProcessTree } from "../../shared/process-tree.ts";
import { saveProgress, cleanupOldRuns } from "./artifacts.ts";

// ── Types ──────────────────────────────────────────────────────────────────────

export interface ActiveRun {
  runId: string;
  progress: WorkflowProgress;
  abortController: AbortController;
  /** PID of the sandbox child process (if any) */
  childPid?: number;
  /** Extension API reference for follow-up messages */
  pi?: any;
}

// ── Run Registry ────────────────────────────────────────────────────────────────

const activeRuns = new Map<string, ActiveRun>();

/**
 * Register a running workflow. Returns a deregister function.
 */
export function registerRun(run: ActiveRun): () => void {
  activeRuns.set(run.runId, run);
  return () => {
    activeRuns.delete(run.runId);
  };
}

/**
 * Get an active run by ID.
 */
export function getActiveRun(runId: string): ActiveRun | undefined {
  return activeRuns.get(runId);
}

/**
 * Get all active runs.
 */
export function getActiveRuns(): ActiveRun[] {
  return Array.from(activeRuns.values());
}

/**
 * Get the count of active (running) workflows.
 */
export function activeRunCount(): number {
  let count = 0;
  for (const run of activeRuns.values()) {
    if (run.progress.status === "running") count++;
  }
  return count;
}

// ── Graceful Shutdown ───────────────────────────────────────────────────────────

/**
 * Cancel all running workflows gracefully.
 *
 * 1. Aborts all in-flight agent calls
 * 2. Kills sandbox child processes (Windows-adapted)
 * 3. Saves final progress snapshots
 */
export async function shutdownAll(): Promise<void> {
  const runs = Array.from(activeRuns.values());

  const killPromises: Promise<void>[] = [];

  for (const run of runs) {
    run.progress.status = "cancelled";
    run.progress.finishedAt = Date.now();

    // Mark pending/running steps as cancelled
    for (const step of run.progress.steps) {
      if (step.status === "pending" || step.status === "running") {
        step.status = "cancelled";
        step.finishedAt = Date.now();
      }
    }

    // Abort in-flight agent calls
    try {
      run.abortController.abort();
    } catch {
      // Best effort
    }

    // Kill sandbox child process
    if (run.childPid) {
      killPromises.push(killProcessTreeWindowsSafe(run.childPid));
    }

    // Save final progress
    try {
      saveProgress(run.runId, run.progress);
    } catch {
      // Best effort
    }

    activeRuns.delete(run.runId);
  }

  // Wait for all kills to complete (bounded)
  const timeout = Promise.resolve().then(() => new Promise<void>((r) => setTimeout(r, 5_000)));
  await Promise.race([
    Promise.allSettled(killPromises),
    timeout,
  ]);
}

/**
 * Full cleanup: shutdown all runs and delete old artifacts.
 */
export async function fullCleanup(): Promise<void> {
  await shutdownAll();
  cleanupOldRuns(0);
}

// ── Windows-Adapted Process Tree Kill ───────────────────────────────────────────

/**
 * Kill a process and its entire tree, adapted for Windows.
 *
 * Windows: taskkill /PID <pid> /T /F
 * POSIX: process.kill(-pid, SIGKILL) with SIGTERM grace period
 */
export async function killProcessTreeWindowsSafe(pid: number): Promise<void> {
  await killProcessTree(pid, { force: true, graceMs: 500 });
}

// ── Background Follow-Up ────────────────────────────────────────────────────────

/**
 * Perform background follow-up after a workflow completes.
 *
 * If a pi ExtensionAPI reference is available, sends exactly ONE followUp
 * message via pi.sendUserMessage. Otherwise, best-effort cleanup only.
 */
export async function backgroundFollowUp(
  runId: string,
  pi?: any,
): Promise<void> {
  // Cleanup old completed runs (keep last 24h)
  try {
    cleanupOldRuns();
  } catch {
    // Non-critical
  }

  // Remove from active registry
  activeRuns.delete(runId);
}