"""
allay_rt.py  —  AllayRT  (Allay Runtime)
=========================================
A lightweight custom ML inference runtime.

Quick start
-----------
    from allay_rt import AllayRT
    import numpy as np

    art = AllayRT()
    model          = art.read_model("model.art")
    compiled_model = art.compile_model(model, "CPU")
    output         = compiled_model.run(img)

API reference
-------------
AllayRT
    .read_model(path)            → ARTModel
    .compile_model(model, device)→ CompiledModel
    .version                     → str

CompiledModel
    .run(inputs, **kwargs)       → np.ndarray  or  dict
    .benchmark(inputs, n=50)     → BenchmarkResult
    .device                      → str
    .model                       → ARTModel
    .input_names                 → list[str]
    .output_names                → list[str]
"""

import time
import numpy as np
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

from AllayRT.model import ARTModel, load_art
from AllayRT.core import run_graph, register_op, OP_TABLE


__version__ = "1.0.0"
__all__ = ["AllayRT", "CompiledModel", "BenchmarkResult"]


# ─────────────────────────────────────────────
#  Supported devices
# ─────────────────────────────────────────────

SUPPORTED_DEVICES = {"CPU"}

# Future hooks — GPU/NPU backends can be registered here
_BACKEND_REGISTRY: Dict[str, Any] = {}


def register_backend(device: str, backend_fn):
    """
    Register a custom execution backend for a device string.

    Parameters
    ----------
    device     : e.g. "GPU", "NPU", "CUDA"
    backend_fn : callable(nodes, tensors, feed, output_names) → dict
    """
    _BACKEND_REGISTRY[device.upper()] = backend_fn
    SUPPORTED_DEVICES.add(device.upper())
    print(f"[AllayRT] Registered backend for device: {device.upper()}")


# ─────────────────────────────────────────────
#  Benchmark result
# ─────────────────────────────────────────────

@dataclass
class BenchmarkResult:
    device:      str
    n_runs:      int
    mean_ms:     float
    median_ms:   float
    min_ms:      float
    max_ms:      float
    std_ms:      float
    throughput:  float  # inferences/sec

    def __str__(self):
        return (
            f"BenchmarkResult(device={self.device}, n={self.n_runs}, "
            f"mean={self.mean_ms:.2f}ms, median={self.median_ms:.2f}ms, "
            f"min={self.min_ms:.2f}ms, max={self.max_ms:.2f}ms, "
            f"std={self.std_ms:.2f}ms, "
            f"throughput={self.throughput:.1f} inf/s)"
        )


# ─────────────────────────────────────────────
#  Compiled model
# ─────────────────────────────────────────────

class CompiledModel:
    """
    A model compiled for a specific device, ready for inference.

    Obtain via:
        compiled = art.compile_model(model, "CPU")
    """

    def __init__(self, model: ARTModel, device: str, verbose: bool = False):
        self._model   = model
        self._device  = device.upper()
        self._verbose = verbose

        # Pre-build IO maps
        self._input_names  = [inp["name"] for inp in model.inputs]
        self._output_names = [out["name"] for out in model.outputs]

        # Resolve backend
        if self._device in _BACKEND_REGISTRY:
            self._backend = _BACKEND_REGISTRY[self._device]
        elif self._device == "CPU":
            self._backend = None   # use built-in run_graph
        else:
            raise ValueError(
                f"[AllayRT] Unsupported device '{device}'. "
                f"Available: {', '.join(SUPPORTED_DEVICES)}"
            )

        print(f"[AllayRT] Compiled '{model.meta.get('name','model')}' "
              f"on {self._device}  "
              f"({len(model.nodes)} nodes, {len(model.tensors)} weight tensors)")

    # ── Properties ──────────────────────────

    @property
    def device(self) -> str:
        return self._device

    @property
    def model(self) -> ARTModel:
        return self._model

    @property
    def input_names(self) -> List[str]:
        return self._input_names

    @property
    def output_names(self) -> List[str]:
        return self._output_names

    # ── Run ─────────────────────────────────

    def run(
        self,
        inputs: Union[np.ndarray, Dict[str, np.ndarray], List[np.ndarray]],
        output_names: Optional[List[str]] = None,
        return_dict: bool = False,
    ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
        """
        Run inference.

        Parameters
        ----------
        inputs       : single array (if one input), list of arrays, or {name: array} dict
        output_names : subset of outputs to return (default: all)
        return_dict  : if True, always return {name: array} dict

        Returns
        -------
        np.ndarray if single output and return_dict=False, else dict
        """
        feed = self._normalize_inputs(inputs)
        out_names = output_names or self._output_names

        if self._backend is not None:
            results = self._backend(
                self._model.nodes,
                self._model.tensors,
                feed,
                out_names,
            )
        else:
            results = run_graph(
                self._model.nodes,
                self._model.tensors,
                feed,
                out_names,
                device=self._device,
                verbose=self._verbose,
            )

        if not return_dict and len(out_names) == 1:
            return results[out_names[0]]
        return results

    def _normalize_inputs(
        self,
        inputs: Union[np.ndarray, Dict, List],
    ) -> Dict[str, np.ndarray]:
        """Accept array, list-of-arrays, or dict; return {name: array}."""
        if isinstance(inputs, dict):
            return {k: np.asarray(v) for k, v in inputs.items()}
        if isinstance(inputs, np.ndarray):
            if len(self._input_names) != 1:
                raise ValueError(
                    f"[AllayRT] Model has {len(self._input_names)} inputs "
                    f"but a single array was provided. "
                    f"Pass a dict: {{{self._input_names[0]!r}: array, ...}}"
                )
            return {self._input_names[0]: inputs}
        if isinstance(inputs, (list, tuple)):
            if len(inputs) != len(self._input_names):
                raise ValueError(
                    f"[AllayRT] Expected {len(self._input_names)} inputs, "
                    f"got {len(inputs)}"
                )
            return {n: np.asarray(v) for n, v in zip(self._input_names, inputs)}
        raise TypeError(
            f"[AllayRT] inputs must be np.ndarray, list, or dict; got {type(inputs)}"
        )

    # ── Benchmark ───────────────────────────

    def benchmark(
        self,
        inputs: Union[np.ndarray, Dict, List],
        n: int = 50,
        warmup: int = 5,
    ) -> BenchmarkResult:
        """
        Benchmark inference latency.

        Parameters
        ----------
        inputs : same as .run()
        n      : number of timed runs
        warmup : number of warmup runs (not counted)

        Returns
        -------
        BenchmarkResult
        """
        print(f"[AllayRT] Benchmarking on {self._device} "
              f"(warmup={warmup}, n={n}) ...")

        for _ in range(warmup):
            self.run(inputs)

        times_ms = []
        for _ in range(n):
            t0 = time.perf_counter()
            self.run(inputs)
            times_ms.append((time.perf_counter() - t0) * 1000)

        times = np.array(times_ms)
        result = BenchmarkResult(
            device     = self._device,
            n_runs     = n,
            mean_ms    = float(times.mean()),
            median_ms  = float(np.median(times)),
            min_ms     = float(times.min()),
            max_ms     = float(times.max()),
            std_ms     = float(times.std()),
            throughput = 1000.0 / float(times.mean()),
        )
        print(f"[AllayRT] {result}")
        return result

    def __repr__(self):
        return (
            f"<CompiledModel device={self._device} "
            f"name={self._model.meta.get('name','?')} "
            f"inputs={self._input_names} outputs={self._output_names}>"
        )


# ─────────────────────────────────────────────
#  AllayRT  (main entry point)
# ─────────────────────────────────────────────

class AllayRT:
    """
    AllayRT — Allay Runtime
    =======================
    A custom ML inference runtime for .art model files.

    Usage
    -----
        art = AllayRT()
        model          = art.read_model("model.art")
        compiled_model = art.compile_model(model, "CPU")
        output         = compiled_model.run(img)

    Custom ops
    ----------
        import numpy as np
        from allay_rt import AllayRT
        from allay_rt_core import register_op

        def my_custom_op(inputs, attrs, ctx):
            return [inputs[0] * 2]

        register_op("MyOp", my_custom_op)
    """

    version = __version__

    def __init__(self, verbose: bool = False):
        self._verbose = verbose
        print(f"[AllayRT] Runtime v{self.version} initialized")

    # ── Model loading ────────────────────────

    def read_model(self, model: str) -> ARTModel:
        """
        Load a .art model file from disk.

        Parameters
        ----------
        model : path to the .art file

        Returns
        -------
        ARTModel
        """
        return load_art(model)

    # ── Compilation ──────────────────────────

    def compile_model(
        self,
        model: ARTModel,
        device: str = "CPU",
    ) -> CompiledModel:
        """
        Compile an ARTModel for inference on the specified device.

        Parameters
        ----------
        model  : ARTModel loaded via read_model()
        device : "CPU" (default) | custom registered device

        Returns
        -------
        CompiledModel  — call .run(inputs) to infer
        """
        return CompiledModel(model, device=device, verbose=self._verbose)

    # ── Utility ──────────────────────────────

    @staticmethod
    def supported_ops() -> List[str]:
        """Return list of all natively supported ONNX ops."""
        return sorted(OP_TABLE.keys())

    @staticmethod
    def register_op(name: str, fn):
        """Register a custom op. See runtime.register_op."""
        register_op(name, fn)

    @staticmethod
    def register_backend(device: str, fn):
        """Register a custom device backend."""
        register_backend(device, fn)

    def __repr__(self):
        return f"<AllayRT v{self.version} verbose={self._verbose}>"
