// Automatically inject a Nix development environment into every shell
// execution (the agent's `bash` tool and Paseo's integrated terminal),
// based on the working directory of each command.
//
// This file is the source of truth for the nix-develop plugin. It is bundled
// into the opencode-chamber wrapper (see ../../flake.nix) and loaded via
// OPENCODE_CONFIG_CONTENT's `plugin` array, so `nix run .` ships it
// automatically. To change the baseline tool list, edit BASELINE_PACKAGES
// below (and bump BASELINE_VERSION) and rebuild.
//
// For each shell call we walk up from `cwd` to find the nearest `flake.nix`.
// We then ask Nix to evaluate a SINGLE dev environment that is the union of:
//
//   * the repo's own `devShells.<system>.default` (when the flake defines one), and
//   * a fixed set of baseline tools that every session should always have
//     (nix, git, ripgrep, findutils, coreutils, ...).
//
// The union is expressed in Nix via `mkShell { inputsFrom = [ repoShell ]; ... }`
// and resolved with `nix print-dev-env --json --expr`. This is a single
// evaluation with a single closure and Nix-computed PATH precedence — the repo
// shell's tools take precedence over baseline, and both are derived (when a
// flake is present) from the repo's OWN nixpkgs input, so there is no version
// skew. When no flake is found we still provide the baseline tools using the
// `nixpkgs` flake registry entry.
//
// The resulting exported variables are merged into the command's environment.
// Path-like variables (PATH, etc.) are prepended to the server's inherited
// values rather than replacing them, so anything already on PATH survives.
//
// Caching: the resolved environment is cached per directory. The cache key is
// a content hash of `flake.nix` + `flake.lock` (plus a version tag for the
// baseline list), so editing the flake or running `nix flake update`
// invalidates the cache automatically with no server restart. Directories with
// no flake share a single baseline-only cache entry.

import { createHash } from "node:crypto"
import { existsSync, readFileSync } from "node:fs"
import { join, dirname, parse, delimiter } from "node:path"

// Baseline tools every session should always have available, as nixpkgs
// attribute names. Edit this list to change the baseline.
const BASELINE_PACKAGES = [
  "nix",
  "git",
  "jujutsu", // provides the `jj` command
  "ripgrep",
  "fd",
  "findutils",
  "coreutils",
  "gnugrep",
  "gnused",
  "gawk",
  "which",
  "util-linux", // provides the `setsid` command
  "curl",
]

// Bump when BASELINE_PACKAGES changes so cached entries are invalidated.
const BASELINE_VERSION = "3"

// Variables that `nix print-dev-env` emits from the stdenv build sandbox but
// that must never leak into an interactive agent/terminal session. The worst
// offender is HOME=/homeless-shelter (the Nix build sandbox home): propagating
// it sends every tool's config/cache/credential lookups into a nonexistent
// directory. The build-scratch and bookkeeping vars below are equally bogus
// outside a derivation. We strip these entirely so the session keeps the
// daemon's real inherited values (real $HOME, real $TMPDIR, ...).
const SANDBOX_VARS = new Set([
  "HOME",
  "TMPDIR",
  "TMP",
  "TEMP",
  "TEMPDIR",
  "PWD",
  "OLDPWD",
  "SHLVL",
  "NIX_BUILD_TOP",
  "NIX_BUILD_CORES",
  "NIX_LOG_FD",
  "NIX_ENFORCE_PURITY",
  "NIX_STORE",
  "TZ",
  "builder",
  "name",
  "out",
  "outputs",
  "system",
  "phases",
  "configureFlags",
  "dontAddDisableDepTrack",
])

// Colon-separated path-like variables. For these we merge the dev shell value
// with the existing (inherited) value instead of replacing it, so baseline and
// system tools stay reachable.
const PATH_LIKE_VARS = new Set([
  "PATH",
  "LD_LIBRARY_PATH",
  "LIBRARY_PATH",
  "PKG_CONFIG_PATH",
  "XDG_DATA_DIRS",
  "MANPATH",
  "INFOPATH",
  "ACLOCAL_PATH",
  "CMAKE_PREFIX_PATH",
  "C_INCLUDE_PATH",
  "CPLUS_INCLUDE_PATH",
])

// Prepend `head` entries before `tail` entries, dropping empties and
// duplicates while preserving first-seen order. Either side may be undefined.
function mergePath(head, tail) {
  const seen = new Set()
  const out = []
  for (const part of `${head ?? ""}${delimiter}${tail ?? ""}`.split(delimiter)) {
    if (part && !seen.has(part)) {
      seen.add(part)
      out.push(part)
    }
  }
  return out.join(delimiter)
}

// Walk up from `start` to find the nearest directory containing flake.nix.
function findFlakeDir(start) {
  let dir = start
  const root = parse(dir).root
  while (true) {
    if (existsSync(join(dir, "flake.nix"))) return dir
    if (dir === root) return null
    dir = dirname(dir)
  }
}

// Cache key for a directory. For a flake dir this is a content hash of
// flake.nix + flake.lock; for a non-flake dir it is a constant. The baseline
// version is always mixed in so changing the baseline invalidates everything.
//
// Note: for a flake with no committed flake.lock yet, the first evaluation
// creates one, which changes this hash and causes exactly one extra evaluation
// on the following call. Flakes with a committed lock are stable immediately.
function cacheKey(flakeDir) {
  const hash = createHash("sha256")
  hash.update(`baseline:${BASELINE_VERSION}\0`)
  if (flakeDir) {
    for (const name of ["flake.nix", "flake.lock"]) {
      const file = join(flakeDir, name)
      hash.update(name)
      hash.update("\0")
      try {
        hash.update(readFileSync(file))
      } catch {
        hash.update("\0__absent__\0")
      }
      hash.update("\0")
    }
  } else {
    hash.update("no-flake\0")
  }
  return hash.digest("hex")
}

// print-dev-env creates a temporary build directory under TMPDIR. If the
// inherited TMPDIR points at a nonexistent path (e.g. a stdenv `/build` leaked
// into the server's environment) the call fails, so fall back to /tmp.
function safeTmpdir() {
  const t = process.env.TMPDIR
  return t && existsSync(t) ? t : "/tmp"
}

// Build the Nix expression that yields the union dev shell.
//
// With a flake: import it, reuse its nixpkgs input, and wrap its default dev
// shell (if any) with the baseline packages via `inputsFrom`.
//
// Without a flake: use the `nixpkgs` registry entry and a baseline-only shell.
function unionExpr(flakeDir) {
  const baseline = BASELINE_PACKAGES.map((p) => `pkgs.${p}`).join(" ")
  if (flakeDir) {
    // JSON.stringify gives us a correctly-quoted Nix string literal for the path.
    const flakeRef = JSON.stringify(flakeDir)
    return `
      let
        system = builtins.currentSystem;
        repo = builtins.getFlake ${flakeRef};
        pkgs = import repo.inputs.nixpkgs { inherit system; };
        realShell = repo.devShells.\${system}.default or null;
        baseline = [ ${baseline} ];
      in {
        devShell = pkgs.mkShell ({ packages = baseline; }
          // (if realShell == null then {} else { inputsFrom = [ realShell ]; }));
      }`
  }
  return `
    let
      system = builtins.currentSystem;
      pkgs = import <nixpkgs> { inherit system; };
      baseline = [ ${baseline} ];
    in {
      devShell = pkgs.mkShell { packages = baseline; };
    }`
}

export const NixDevelopPlugin = async ({ $, client, directory }) => {
  const log = (level, message, extra) =>
    client.app
      .log({ body: { service: "nix-develop", level, message, extra } })
      .catch(() => {})

  // key string -> { env, ok, promise? }
  const cache = new Map()

  // sessionID -> last cacheKey() we posted a notice for. Lets us tell the
  // agent, once per session, whether a flake dev shell is active — and again
  // if flake.nix/flake.lock changes mid-session — without adding a separate
  // chat message. We piggyback on the first `bash` tool result instead, since
  // the agent already reads that.
  const announced = new Map()

  function noticeFor(flakeDir, ok, isReload) {
    const verb = isReload ? "reloaded" : "loaded"
    if (flakeDir) {
      const status = ok
        ? `flake dev shell ${verb} from ${flakeDir} (merged with baseline tools: git, rg, fd, curl, ...)`
        : `found flake.nix at ${flakeDir} but evaluating its dev shell failed; baseline tools only (see nix-develop logs)`
      return `[nix-develop] ${status}. No need to run \`nix develop\`/\`nix shell\` yourself — it's already applied to every bash call.`
    }
    return `[nix-develop] no flake.nix found; baseline tools only (git, rg, fd, curl, ...). Missing tools? Run \`nix shell nixpkgs#<pkg> -c <cmd>\`.`
  }

  async function evaluate(flakeDir) {
    const where = flakeDir ?? "(no flake)"
    try {
      const expr = unionExpr(flakeDir)
      // `--impure` is required for builtins.getFlake / <nixpkgs>. The expression
      // is passed as a single interpolated argument; Bun Shell escapes it and
      // does not invoke a system shell, so there are no quoting/injection risks.
      const res = await $`nix print-dev-env --json --impure --expr ${expr} devShell`
        .env({ ...process.env, TMPDIR: safeTmpdir() })
        .quiet()
      const parsed = JSON.parse(res.stdout.toString())
      const env = {}
      for (const [k, v] of Object.entries(parsed.variables ?? {})) {
        // Only plain exported scalars. print-dev-env also emits bash functions
        // and array-typed vars, which cannot be represented as simple env vars.
        if (v && v.type === "exported" && typeof v.value === "string") {
          env[k] = v.value
        }
      }
      log("info", `resolved nix dev env for ${where}`, {
        vars: Object.keys(env).length,
      })
      return { env, ok: true }
    } catch (e) {
      log("warn", `nix print-dev-env failed for ${where}`, {
        error: String(e?.stderr ?? e),
      })
      // Cache an empty env for this key so we don't retry on every command;
      // a change to flake.nix/flake.lock (or the baseline) will retry.
      return { env: {}, ok: false }
    }
  }

  async function resolveEnv(flakeDir) {
    const key = cacheKey(flakeDir)
    const entry = cache.get(key)
    if (entry) return entry.promise ?? entry

    const promise = evaluate(flakeDir)
    cache.set(key, { promise })
    const result = await promise
    cache.set(key, result)
    return result
  }

  return {
    "shell.env": async (input, output) => {
      // A flake dir gives the repo shell + baseline; no flake gives baseline only.
      const flakeDir = findFlakeDir(input.cwd)
      const { env } = await resolveEnv(flakeDir)

      for (const [k, v] of Object.entries(env)) {
        if (SANDBOX_VARS.has(k)) {
          // Build-sandbox leakage (HOME=/homeless-shelter, build scratch dirs,
          // derivation bookkeeping). Never propagate these into the session.
          continue
        }
        if (PATH_LIKE_VARS.has(k)) {
          // Prepend dev-shell entries ahead of the inherited value, so the dev
          // shell + baseline take precedence while anything already on PATH
          // (e.g. system binaries) survives.
          output.env[k] = mergePath(v, output.env[k])
        } else if (output.env[k] === undefined) {
          // Non-path var: let the dev env provide it, but never clobber a value
          // the caller set intentionally.
          output.env[k] = v
        }
      }
    },

    // Tell the agent, once per session, whether a flake dev shell is active —
    // and again if flake.nix/flake.lock changes — by prepending a short note
    // to the first `bash` tool result. This rides on output the agent already
    // reads instead of adding a separate chat message, and stays silent on
    // every call after that until the flake actually changes.
    "tool.execute.after": async (input, output) => {
      if (input.tool !== "bash") return
      if (!output || typeof output.output !== "string") return

      const sessionID = input.sessionID ?? "global"
      const flakeDir = findFlakeDir(input.args?.workdir ?? directory)
      const key = cacheKey(flakeDir)

      const prevKey = announced.get(sessionID)
      if (prevKey === key) return
      announced.set(sessionID, key)

      const entry = cache.get(key)
      const ok = entry && !entry.promise ? entry.ok : true
      const notice = noticeFor(flakeDir, ok, prevKey !== undefined)
      output.output = `${notice}\n\n${output.output}`
    },
  }
}