"""direnv for the model's jobs: before each `run()`, the kernel's environment
catches up with the .envrc direnv would load for the kernel's folder, as
direnv's shell hook does before each prompt (robot-docs/direnv.md). The kernel
is the shell: os.chdir into another project loads that one's .envrc, and
leaving it unloads it.

direnv decides what loads: a file the person allowed loads, a blocked or
denied one does not. `direnv export json` answers only what changed since the
last call (the DIRENV_* variables it leaves in the environment say what it
loaded), so an unchanged .envrc costs a few milliseconds and an edited one
reloads by itself. Jobs inherit os.environ, so every program the kernel starts
afterwards sees the same environment."""

from __future__ import annotations

import json
import os
import re
import shutil
import subprocess
import tempfile
from pathlib import Path

import albedo_trace

# seconds a first evaluation may take: a `use flake` .envrc builds a dev shell
TIMEOUT = 600
COLOUR = re.compile(r"\x1b\[[0-9;]*m")
BLOCKED = re.compile(r"direnv: error (.*) is blocked\..*")

# the direnv this kernel runs, found once; "" when there is none
_direnv: str | None = None
# whether this kernel loaded an environment it may have to take back; one the
# daemon inherited from the shell that started it stays alone
_loaded = False


def refresh() -> None:
    """Load, reload or unload the .envrc for the kernel's folder, printing what
    direnv says when something changed. Never raises: a job started without
    the environment beats no job."""
    global _direnv, _loaded
    if _direnv is None:
        _direnv = shutil.which("direnv") or ""
    here = os.getcwd()
    # nothing to load here, and nothing loaded to take back
    if not _direnv or not (_loaded or _has_envrc(Path(here))):
        return
    try:
        with albedo_trace.unobserved():
            done = subprocess.run(
                [_direnv, "export", "json"],
                cwd=here,
                env=dict(os.environ),
                capture_output=True,
                text=True,
                timeout=TIMEOUT,
            )
    except subprocess.TimeoutExpired:
        print(
            f"direnv: gave up loading .envrc for {here} after {TIMEOUT} s; jobs start without it"
        )
        return
    except OSError as error:
        print(f"direnv: could not run {_direnv}: {error}")
        return
    stderr = COLOUR.sub("", done.stderr)
    blocked = any(BLOCKED.fullmatch(line.strip()) for line in stderr.splitlines())
    if blocked and _loaded:
        # A changed .envrc is blocked, but direnv's JSON still describes the
        # old environment. Ask direnv to leave it from a directory with no
        # .envrc before applying that stale export.
        try:
            with tempfile.TemporaryDirectory() as empty:
                unload = subprocess.run(
                    [_direnv, "export", "json"],
                    cwd=empty,
                    env=dict(os.environ),
                    capture_output=True,
                    text=True,
                    timeout=TIMEOUT,
                )
            change = json.loads(unload.stdout) if unload.stdout.strip() else {}
        except (OSError, subprocess.TimeoutExpired, json.JSONDecodeError):
            change = {}
    else:
        try:
            change = json.loads(done.stdout) if done.stdout.strip() else {}
        except json.JSONDecodeError:
            change = {}
    for name, value in change.items() if isinstance(change, dict) else ():
        if value is None:
            os.environ.pop(name, None)
        elif isinstance(value, str):
            os.environ[name] = value
    _loaded = "DIRENV_DIR" in os.environ
    for line in stderr.splitlines():
        blocked = BLOCKED.fullmatch(line.strip())
        if blocked:
            # the person allows it, not the model
            line = f"direnv: {blocked[1]} is not allowed, so jobs start without it until the person allows it"
        if line.strip():
            print(line)


def _has_envrc(directory: Path) -> bool:
    return any(
        (folder / ".envrc").is_file() for folder in (directory, *directory.parents)
    )