"""Load agent.py or decide.py for a check, with the real library or with a stand-in for it.

Both files run the model themselves: they start with `from decider.infer import Decider`.
That library is big (it brings PyTorch) and the model is a 3.8 GB download. So:

- where decider-ai is installed, the checks use it, and load the model only once;
- where it is not, a stand-in takes its place. It has the one method the two files call,
  system_one(). It sends the same state and the same questions to the model's server,
  which runs the same model. That is the address in ask.py, or the one you give.

load_agent() says which of the two it used, and the checks print it.
"""
import json
import pathlib
import sys
import types
import urllib.request

ASSETS = pathlib.Path(__file__).resolve().parent.parent      # the tutorial's files: one folder up from checks/
LOADED = {}               # the real model, kept after its first load: loading is slow
REAL = {}                 # the library's own Decider class, kept before anything replaces it


class ServerDecider:
    """The stand-in for the library's Decider: the same question, sent to the model's server."""
    address = None

    def __init__(self, name, **options):
        self.name = name

    def system_one(self, state, questions):
        message = json.dumps({"state": state, "questions": questions}).encode()
        request = urllib.request.Request(self.address, message, {"Content-Type": "application/json"})
        with urllib.request.urlopen(request, timeout=30) as reply:
            return json.load(reply)


def server_address(address):
    if address:
        return address
    sys.path.insert(0, str(ASSETS))
    import ask
    return ask.MODEL


def load_once(name, **options):
    """The real Decider, but each model is loaded only once."""
    if name not in LOADED:
        LOADED[name] = REAL["Decider"](name, **options)
    return LOADED[name]


def run_with_library(source, file_name, names):
    import decider.infer
    REAL.setdefault("Decider", decider.infer.Decider)
    decider.infer.Decider = load_once        # so the file's own import gets load_once
    try:
        exec(compile(source, file_name, "exec"), names)
    finally:
        decider.infer.Decider = REAL["Decider"]


def run_with_stand_in(source, file_name, names, address):
    ServerDecider.address = server_address(address)
    fake = types.ModuleType("decider.infer")
    fake.Decider = ServerDecider
    before = {module: sys.modules.get(module) for module in ["decider", "decider.infer"]}
    sys.modules["decider"] = types.ModuleType("decider")
    sys.modules["decider.infer"] = fake
    try:
        exec(compile(source, file_name, "exec"), names)
    finally:
        for module, kept in before.items():
            if kept is None:
                sys.modules.pop(module, None)
            else:
                sys.modules[module] = kept


def load_agent(source=None, address=None, stand_in=None, name="agent.py"):
    """Run the file's text, changed or not, and give back its names and "library" or "server".

    agent.py only defines things: call names["main"]() to run it. decide.py has no main():
    running its text asks its one question and prints the answer.
    stand_in=True uses the stand-in even where the library is installed."""
    if source is None:
        source = (ASSETS / name).read_text(encoding="utf-8")
    if stand_in is None:
        try:
            import decider.infer
            stand_in = False
        except ImportError:
            stand_in = True
    names = {"__name__": "checked_" + name.replace(".py", "")}
    if stand_in:
        run_with_stand_in(source, name, names, address)
        return names, "server"
    run_with_library(source, name, names)
    return names, "library"