"""Record modules a ready daemon has loaded, for standbys, and those a first
text turn loads after that, 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/ready-modules.txt and 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")
    # A standby would have loaded more than the boot itself does.
    daemon.env.update(ALBEDO_INSPECT="1", ALBEDO_PRELOAD="0", ALBEDO_STANDBY="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 before, 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")
    parser.add_argument("--ready", type=Path, default=ROOT / "priv/ready-modules.txt")
    args = parser.parse_args()
    harness.snapshot_daemon()
    chat_ready, chat_turn = record("chat_completions")
    responses_ready, responses_turn = record("responses")
    for path, modules in (
        (args.ready, chat_ready | responses_ready),
        (args.output, chat_turn | responses_turn),
    ):
        path.write_text("\n".join(sorted(modules)) + "\n")
        print(f"{len(modules)} modules written to {path}")


if __name__ == "__main__":
    main()