"""Record modules a first text turn loads for the daemon's background code warm.

Run with nix develop -c python3 test/manual/record_boot_modules.py. Uses a private
fake-provider daemon, no credentials or Python kernel; writes priv/boot-modules.txt.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys
import subprocess
from types import SimpleNamespace

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "test/e2e"))
import harness  # noqa: E402
from harness import Albedo, Daemon, Provider, text  # noqa: E402
from inspect_support import Inspection  # noqa: E402


def loaded(app):
    return set(
        json.loads(
            Inspection(app).evaluate(
                "io:put_chars(json:encode([atom_to_binary(M) || {M,_} <- Call(code, all_loaded, [])]))"
            )
        )
    )


def record(protocol):
    daemon = Daemon("boot-modules")
    daemon.env.update(ALBEDO_INSPECT="1", ALBEDO_PRELOAD="0")
    try:
        daemon.boot()
        before = loaded(
            SimpleNamespace(
                daemon=daemon, root=daemon.root, home=daemon.home, env=daemon.env
            )
        )
        with (
            harness.on_daemon(daemon),
            Albedo(Provider(lambda _: text("ok")), protocol=protocol) as app,
        ):
            session = app.session()
            app.prompt(session, "hello").close()
            app.idle(session)
            return loaded(app) - before
    finally:
        daemon.shutdown()
        remaining = [
            line
            for line in subprocess.check_output(
                ["ps", "-axo", "pid=,command="], text=True
            ).splitlines()
            if str(daemon.home) in line
        ]
        if remaining:
            raise AssertionError(remaining)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=ROOT / "priv/boot-modules.txt")
    args = parser.parse_args()
    harness.snapshot_daemon()
    modules = sorted(record("chat_completions") | record("responses"))
    args.output.write_text("\n".join(modules) + "\n")
    print(f"{len(modules)} modules written to {args.output}")


if __name__ == "__main__":
    main()