import { chmod, mkdir, readFile, writeFile } from "node:fs/promises";
import { hostname } from "node:os";
import { join, resolve } from "node:path";
import { ControlPlaneError, WorkerClient } from "@autoscholar/maclustr-adapter";
import { createLogger, redact } from "@autoscholar/observability";
import { type RunnerHost, runResearch, simulateResearch } from "@autoscholar/research-runner";
import type { ArtifactKind, RunEventInput, StageKey } from "@autoscholar/schemas";
import { collectNodeMetadata, loadAverage, sdkVersion, workerVersion } from "./node-meta.js";
import { runSelfTest } from "./selftest.js";

const log = createLogger("worker");

interface Config {
  controlPlaneUrl: string;
  workerName: string;
  workerRoot: string;
  token?: string;
  enrollmentSecret?: string;
  mode: "real" | "simulated";
  maxConcurrent: number;
  profiles: string[];
  apiKey?: string;
  extraDomains: string[];
  tectonicBin?: string;
  latexTemplateDir: string;
  pollMs: number;
}

function loadConfig(): Config {
  const e = process.env;
  const workerRoot = e.WORKER_ROOT ?? "/var/autoscholar";
  return {
    controlPlaneUrl: e.CONTROL_PLANE_URL ?? "http://127.0.0.1:8470",
    workerName: e.WORKER_NAME ?? hostname(),
    workerRoot,
    token: e.WORKER_TOKEN || undefined,
    enrollmentSecret: e.WORKER_ENROLLMENT_SECRET || undefined,
    mode: e.WORKER_MODE === "simulated" ? "simulated" : "real",
    maxConcurrent: Number(e.WORKER_MAX_CONCURRENT ?? 1),
    profiles: (e.WORKER_PROFILES ?? "standard")
      .split(",")
      .map((s) => s.trim())
      .filter(Boolean),
    apiKey: e.ANTHROPIC_API_KEY || undefined,
    extraDomains: (e.RESEARCH_ALLOWED_DOMAINS ?? "")
      .split(",")
      .map((s) => s.trim())
      .filter(Boolean),
    tectonicBin: e.TECTONIC_BIN || undefined,
    latexTemplateDir: e.LATEX_TEMPLATE_DIR ?? resolve(workerRoot, "templates/latex"),
    pollMs: Number(e.WORKER_POLL_MS ?? 4000),
  };
}

class Worker {
  private readonly client: WorkerClient;
  private readonly running = new Map<
    string,
    { fencingToken: number; cancel: boolean; restart?: boolean; events: RunEventInput[] }
  >();
  private stopping = false;
  private draining = false;
  private tokenPath: string;

  constructor(private readonly cfg: Config) {
    this.client = new WorkerClient({ baseUrl: cfg.controlPlaneUrl, token: cfg.token });
    this.tokenPath = join(cfg.workerRoot, "state/worker-token");
  }

  async start() {
    await mkdir(join(this.cfg.workerRoot, "runs"), { recursive: true });
    await mkdir(join(this.cfg.workerRoot, "state"), { recursive: true });
    if (!this.cfg.token) {
      try {
        const saved = (await readFile(this.tokenPath, "utf8")).trim();
        if (saved) this.client.setToken(saved);
      } catch {
        /* none */
      }
    }
    const node = collectNodeMetadata({
      simulated: this.cfg.mode === "simulated",
      profiles: this.cfg.profiles,
    });
    if (this.cfg.mode === "real" && !node.sandboxAvailable)
      log.error(
        "sandbox-exec not available: this worker will never be dispatched research (sandboxAvailable=false)",
      );
    if (this.cfg.mode === "real" && !this.cfg.apiKey)
      log.error("ANTHROPIC_API_KEY missing: research sessions cannot start");
    // enroll if needed
    let enrolled = false;
    for (let i = 0; i < 30 && !enrolled && !this.stopping; i++) {
      try {
        await this.client.heartbeat({
          status: "online",
          capacity: this.cfg.maxConcurrent,
          running: [],
          load: loadAverage(),
          node,
          leases: [],
        });
        enrolled = true;
      } catch (e) {
        if (e instanceof ControlPlaneError && e.status === 401 && this.cfg.enrollmentSecret) {
          try {
            const r = await this.client.enroll(this.cfg.enrollmentSecret, this.cfg.workerName, node);
            this.client.setToken(r.token);
            await writeFile(this.tokenPath, r.token, { mode: 0o600 });
            await chmod(this.tokenPath, 0o600);
            log.info({ workerId: r.workerId }, "enrolled with control plane");
            enrolled = true;
          } catch (e2) {
            log.error({ err: (e2 as Error).message }, "enrollment failed");
          }
        } else log.warn({ err: (e as Error).message, attempt: i }, "control plane unreachable; retrying");
        if (!enrolled) await new Promise((r) => setTimeout(r, 5000));
      }
    }
    if (!enrolled) throw new Error("could not reach/enroll with the control plane");
    log.info(
      {
        name: this.cfg.workerName,
        mode: this.cfg.mode,
        sdk: sdkVersion(),
        worker: workerVersion(),
        controlPlane: this.cfg.controlPlaneUrl,
        sandbox: node.sandboxAvailable,
      },
      "worker online",
    );
    void this.heartbeatLoop(node);
    void this.leaseLoop();
  }

  private async heartbeatLoop(node: ReturnType<typeof collectNodeMetadata>) {
    while (!this.stopping) {
      try {
        const leases = [...this.running.entries()].map(([runId, r]) => ({
          runId,
          fencingToken: r.fencingToken,
        }));
        const hb = await this.client.heartbeat({
          status: this.draining ? "draining" : "online",
          capacity: this.cfg.maxConcurrent,
          running: [...this.running.keys()],
          load: loadAverage(),
          leases,
          node: { tools: node.tools, sandboxAvailable: node.sandboxAvailable },
        });
        this.draining = hb.drain;
        for (const id of hb.cancel) {
          const r = this.running.get(id);
          if (r) r.cancel = true;
        }
        for (const id of hb.revokedLeases) {
          const r = this.running.get(id);
          if (r) {
            log.warn({ runId: id }, "lease revoked by control plane (fencing) — aborting local work");
            r.cancel = true;
          }
        }
        await this.flushEvents();
      } catch (e) {
        log.warn({ err: (e as Error).message }, "heartbeat failed");
      }
      await new Promise((r) => setTimeout(r, 20_000));
    }
  }

  private async flushEvents() {
    for (const [runId, r] of this.running) {
      if (r.events.length === 0) continue;
      const batch = r.events.splice(0, 200);
      try {
        await this.client.events(
          runId,
          r.fencingToken,
          batch.map((e) => ({ ...e, level: e.level ?? "info" })),
        );
      } catch (e) {
        if (e instanceof ControlPlaneError && e.status === 409) {
          r.cancel = true;
          log.warn({ runId }, "lease lost while posting events; cancelling local work");
        } else r.events.unshift(...batch.slice(0, 50));
      }
    }
  }

  private async leaseLoop() {
    while (!this.stopping) {
      if (!this.draining && this.running.size < this.cfg.maxConcurrent) {
        try {
          const { job } = await this.client.lease(this.cfg.profiles);
          if (job) {
            log.info({ runId: job.runId, attempt: job.attemptNo, resume: job.resumeFromStage }, "leased job");
            void this.execute(job);
          }
        } catch (e) {
          log.warn({ err: (e as Error).message }, "lease request failed");
        }
      }
      await new Promise((r) => setTimeout(r, this.cfg.pollMs));
    }
  }

  private async execute(job: Awaited<ReturnType<WorkerClient["lease"]>>["job"] & object) {
    const state = {
      fencingToken: job.fencingToken,
      cancel: job.cancelRequested,
      events: [] as RunEventInput[],
    };
    this.running.set(job.runId, state);
    const renew = setInterval(async () => {
      try {
        const r = await this.client.renew(job.runId, job.fencingToken);
        if (r.cancelRequested) state.cancel = true;
        await this.flushEvents();
      } catch (e) {
        if (e instanceof ControlPlaneError && e.status === 409) {
          state.cancel = true;
          log.warn({ runId: job.runId }, "lease lost (stale fencing token); aborting");
        }
      }
    }, 30_000);
    const host: RunnerHost = {
      workspaceRoot: this.cfg.workerRoot,
      apiKey: this.cfg.apiKey ?? "",
      extraDomains: this.cfg.extraDomains,
      tectonicBin: this.cfg.tectonicBin,
      tectonicCache: join(this.cfg.workerRoot, "cache/tectonic"),
      latexTemplateDir: this.cfg.latexTemplateDir,
      workerName: this.cfg.workerName,
      sdkVersion: sdkVersion(),
      toolchain: collectNodeMetadata({
        simulated: this.cfg.mode === "simulated",
        profiles: this.cfg.profiles,
      }).tools,
      emit: (ev) => {
        state.events.push(redact(ev));
        if (state.events.length >= 50) void this.flushEvents();
      },
      stage: async (report) => {
        await this.flushEvents();
        await this.client.stage(job.runId, { fencingToken: job.fencingToken, ...report });
      },
      upload: async (path: string, kind: ArtifactKind, contentType: string, buf: Buffer) => {
        await this.client.uploadArtifact(
          job.runId,
          job.fencingToken,
          {
            path,
            kind,
            contentType,
            stage: (job.resumeFromStage ?? "preflight") as StageKey,
            synthetic: this.cfg.mode === "simulated",
          },
          buf,
        );
      },
      isCancelled: () => state.cancel || this.stopping,
    };
    let outcome: Awaited<ReturnType<typeof runResearch>>;
    try {
      const runnerJob = {
        runId: job.runId,
        attemptNo: job.attemptNo,
        manifest: job.manifest,
        instructionText: job.instructionText,
        resumeFromStage: job.resumeFromStage,
        checkpoint: job.checkpoint,
      };
      if (this.cfg.mode === "simulated") {
        const sim = process.env.WORKER_SIM_BEHAVIOUR ?? "";
        outcome = await simulateResearch(runnerJob, host, {
          failAtStage: sim.startsWith("fail:") ? sim.slice(5) : undefined,
          hang: sim === "hang",
          delayMs: Number(process.env.WORKER_SIM_DELAY_MS ?? 0),
        });
      } else {
        if (!this.cfg.apiKey) throw new Error("ANTHROPIC_API_KEY not configured on this worker");
        outcome = await runResearch(runnerJob, host);
      }
      if (state.restart && outcome.outcome === "cancelled") {
        outcome = {
          ...outcome,
          outcome: "failed",
          error: "WORKER_RESTART: worker daemon restarted; resume from checkpoint",
        };
      }
    } catch (e) {
      outcome = {
        outcome: "failed",
        error: `runner crashed: ${(e as Error).message}`.slice(0, 2000),
        costUsd: 0,
        turns: 0,
        stagesCompleted: [],
        resourceUsage: { wallSeconds: 0, diskBytes: 0, peakRssBytes: 0 },
      };
      log.error(
        { runId: job.runId, err: (e as Error).message, stack: (e as Error).stack?.split("\n").slice(0, 5) },
        "runner crashed",
      );
    }
    clearInterval(renew);
    await this.flushEvents();
    try {
      await this.client.complete(job.runId, {
        fencingToken: job.fencingToken,
        outcome: outcome.outcome,
        error: outcome.error,
        costUsd: outcome.costUsd,
        summary: { stagesCompleted: outcome.stagesCompleted, turns: outcome.turns },
        resourceUsage: outcome.resourceUsage,
      });
      log.info({ runId: job.runId, outcome: outcome.outcome, costUsd: outcome.costUsd }, "attempt completed");
    } catch (e) {
      log.error({ runId: job.runId, err: (e as Error).message }, "could not report completion (lease lost?)");
    }
    this.running.delete(job.runId);
  }

  async stop() {
    this.stopping = true;
    // graceful restart: abort sessions (isCancelled() → true via `stopping`) but report a retryable failure, not a cancel
    for (const r of this.running.values()) r.restart = true;
    const t0 = Date.now();
    while (this.running.size > 0 && Date.now() - t0 < 60_000) await new Promise((r) => setTimeout(r, 500));
  }
}

const cfg = loadConfig();
if (process.argv[2] === "selftest") {
  runSelfTest(cfg.workerRoot, cfg.apiKey, cfg.latexTemplateDir, cfg.tectonicBin)
    .then((r) => {
      console.log(JSON.stringify(r, null, 2));
      process.exit(r.ok ? 0 : 1);
    })
    .catch((e) => {
      console.error(e);
      process.exit(1);
    });
} else {
  const w = new Worker(cfg);
  process.on("SIGTERM", () => void w.stop().then(() => process.exit(0)));
  process.on("SIGINT", () => void w.stop().then(() => process.exit(0)));
  w.start().catch((e) => {
    log.fatal({ err: (e as Error).message }, "worker failed to start");
    process.exit(1);
  });
}