#!/usr/bin/env python3
"""Test native afk storage/acknowledgments with ASan/UBSan, without the VM.

Uses deps/json/single_include/nlohmann/json.hpp if restored, or --json-header
pointing to the single header at the revision in the plugin's upstream.json.
"""
import argparse
import os
from pathlib import Path
import re
import shutil
import subprocess
import tempfile

ROOT = Path(__file__).resolve().parents[1]
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--json-header", help="Header path; {game} expands to t5/t6 for their respective pinned headers")
args = parser.parse_args()

with tempfile.TemporaryDirectory(prefix="tws-afk-") as temp:
    directory = Path(temp)
    for game in ("t5", "t6"):
        plugin = ROOT / game / "plugins" / f"tws-{game}-gsc-utils"
        header = Path(args.json_header.format(game=game)) if args.json_header else plugin / "deps/json/single_include/nlohmann/json.hpp"
        if not header.is_file():
            parser.error(f"Missing {header}; restore dependencies or pass --json-header")
        shutil.copyfile(header, directory / "json.hpp")
        (directory / "stdinc.hpp").write_text("#include <algorithm>\n#include <fstream>\n#include <sstream>\n#include <json.hpp>\n")
        executable = directory / game
        subprocess.run([os.environ.get("CXX", "c++"), "-std=c++20", "-Wall", "-Wextra", "-Werror",
                        "-fsanitize=address,undefined", "-I", str(directory), "-I", str(plugin / "src/utils"),
                        str(plugin / "src/utils/afk.cpp"), str(ROOT / "tools/tests/afk_cases.cpp"),
                        "-o", str(executable)], check=True)
        work = directory / f"{game}-data"
        subprocess.run([str(executable), str(work)], check=True)
        # Separate OS processes prove snapshots/identities don't depend on VM or
        # DLL memory. A newly booted producer recovers before its own start.
        subprocess.run([str(executable), str(work), "produce"], check=True)
        subprocess.run([str(executable), str(work), "recover"], check=True)
        subprocess.run([str(executable), str(work), "empty"], check=True)

        # Compile the actual registration adapter as well, substituting only
        # engine interfaces. This does not validate the Windows ABI/VM linker.
        stub = (ROOT / "tools/tests/http_component_stub.hpp").read_text()
        stub = stub.replace("std::vector<script_value> get_keys() const", "void emplace_back(const script_value& value) { (*data)[std::to_string(data->size())] = value; }\n        std::vector<script_value> get_keys() const")
        stub = stub.replace('template <typename F> void add(const char*, F) {}', r"""
    inline std::map<std::string, std::function<scripting::script_value(const std::vector<std::string>&)>> functions;
    template <typename F> void add(const char* name, F callback)
    {
        functions[name] = [callback](const std::vector<std::string>& args) -> scripting::script_value
        {
            if constexpr (std::is_invocable_v<F>) return callback();
            else if constexpr (std::is_invocable_v<F, std::string>) return callback(args.at(0));
            else if constexpr (std::is_invocable_v<F, std::string, std::string>) return callback(args.at(0), args.at(1));
            else if constexpr (std::is_invocable_v<F, std::string, std::string, std::string>) return callback(args.at(0), args.at(1), args.at(2));
            else if constexpr (std::is_invocable_v<F, std::string, std::string, std::string, int>) return callback(args.at(0), args.at(1), args.at(2), std::stoi(args.at(3)));
            else return callback(args.at(0), args.at(1), args.at(2), args.at(3));
        };
    }
""")
        (directory / "component_stub.hpp").write_text(stub)
        component = re.sub(r'^#include[^\n]*\n', '', (plugin / "src/component/afk.cpp").read_text(), flags=re.MULTILINE)
        source = directory / "component.cpp"
        source.write_text('#include "component_stub.hpp"\n#include "afk.hpp"\n#include <array>\n' + component + (ROOT / 'tools/tests/afk_component_cases.cpp').read_text())
        subprocess.run([os.environ.get("CXX", "c++"), "-std=c++20", "-Wall", "-Wextra", "-Werror", "-fsanitize=address,undefined",
                        "-I", str(directory), "-I", str(plugin / "src/utils"), str(source),
                        str(plugin / "src/utils/afk.cpp"), "-o", str(directory / "component")], check=True)
        result = subprocess.run([str(directory / "component")], cwd=directory, check=True, capture_output=True, text=True)
        lines = result.stdout.splitlines()
        for operation in ("load", "write", "remove"):
            assert sum(f"operation={operation} outcome=storage_failure" in line for line in lines) == 1, lines
            assert sum(f"operation={operation} outcome=recovered suppressed=1" in line for line in lines) == 1, lines
        assert "private-secret" not in result.stdout and "Authorization" not in result.stdout
        print(f"PASS {game}: native storage, acknowledgment, process restart and component checks", flush=True)