"""PyMC runner — executes compiled module code in the host environment.

Supports three modes:
    runner.run(data)    — Mode 1/2: run inference, return trace
    runner.build(data)  — Mode 3: return pm.Model for custom extension
    runner.code         — raw code string

Host functions can be injected as keyword arguments:
    runner.run(data, hazard=my_hazard_fn)
    runner.build(data, encode=my_encoder)
"""


class PymcRunner:
    def __init__(self, result: dict):
        self.code = result["code"]
        self.data_names = result["data_names"]
        self.param_names = result["param_names"]
        self.inference_mode = result["inference_mode"]

    def _exec_module(self, ext=None):
        """Execute the generated code and return the namespace.

        Args:
            ext: dict of external callables to inject into the module namespace
        """
        ns = dict(ext or {})
        exec(self.code, ns)
        return ns

    def run(self, data=None, **kwargs):
        """Run inference (Mode 1: no data, Mode 2: with data).

        Args:
            data: dict mapping data names to numpy arrays
            **kwargs: callables are injected as host functions,
                      other kwargs passed to pm.sample (draws, tune, chains, target_accept)

        Returns:
            InferenceData trace object
        """
        ext = {k: v for k, v in kwargs.items() if callable(v)}
        opts = {k: v for k, v in kwargs.items() if not callable(v)}
        ns = self._exec_module(ext)
        return ns["run_inference"](data or {}, **opts)

    def build(self, data=None, **kwargs):
        """Build model without running inference (Mode 3).

        Returns a pm.Model for custom extension.

        Args:
            data: dict mapping data names to numpy arrays
            **kwargs: callables are injected as host functions
        """
        ext = {k: v for k, v in kwargs.items() if callable(v)}
        ns = self._exec_module(ext)
        return ns["create_model"](data or {})

    def __repr__(self):
        return (
            f"PymcRunner(params={self.param_names}, "
            f"data={self.data_names}, mode={self.inference_mode})"
        )
