# coding: utf-8
"""axengine / onnxruntime 推理会话封装。"""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np


@dataclass(frozen=True)
class TensorInfo:
    name: str
    shape: tuple[int, ...]
    dtype: np.dtype


def _numpy_dtype(value: Any) -> np.dtype:
    text = str(value).lower()
    mapping = (
        (("tensor(float)", "float32", "fp32", "f32"), np.float32),
        (("tensor(float16)", "float16", "fp16", "f16"), np.float16),
        (("tensor(int64)", "int64", "s64"), np.int64),
        (("tensor(int32)", "int32", "s32"), np.int32),
        (("tensor(uint16)", "uint16", "u16"), np.uint16),
        (("tensor(uint8)", "uint8", "u8"), np.uint8),
    )
    for aliases, dtype in mapping:
        if any(alias in text for alias in aliases):
            return np.dtype(dtype)
    raise ValueError(f"unsupported runtime tensor dtype: {value}")


def _tensor_info(value: Any) -> TensorInfo:
    shape = getattr(value, "shape", None)
    if shape is None:
        shape = getattr(value, "dims", None)
    if shape is None or any(dim is None for dim in shape):
        raise ValueError(f"dynamic or missing tensor shape for {value.name}: {shape}")
    dtype = getattr(value, "dtype", None)
    if dtype is None:
        dtype = getattr(value, "type", None)
    return TensorInfo(
        name=value.name,
        shape=tuple(int(dim) for dim in shape),
        dtype=_numpy_dtype(dtype),
    )


class InferenceSession:
    """.axmodel -> axengine；.onnx -> onnxruntime CPU。"""

    def __init__(self, model_path: str | Path, backend: str | None = None):
        self.path = Path(model_path)
        if not self.path.is_file():
            raise FileNotFoundError(self.path)
        if backend is None:
            backend = "axengine" if self.path.suffix == ".axmodel" else "onnx"
        self.backend = backend
        if backend == "axengine":
            try:
                import axengine
            except ImportError as error:
                raise RuntimeError(
                    "axengine is unavailable; run this backend on an AXERA board"
                ) from error
            self._session = axengine.InferenceSession(str(self.path))
        elif backend == "onnx":
            try:
                import onnxruntime as ort
            except ImportError as error:
                raise RuntimeError("onnxruntime is required for --backend onnx") from error
            options = ort.SessionOptions()
            options.inter_op_num_threads = 1
            options.intra_op_num_threads = 1
            self._session = ort.InferenceSession(
                str(self.path),
                sess_options=options,
                providers=["CPUExecutionProvider"],
            )
        else:
            raise ValueError(f"unsupported backend: {backend}")

        self.inputs = [_tensor_info(value) for value in self._session.get_inputs()]
        self.outputs = [_tensor_info(value) for value in self._session.get_outputs()]
        self.input_by_name = {value.name: value for value in self.inputs}

    def run(self, feed: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
        missing = [value.name for value in self.inputs if value.name not in feed]
        if missing:
            raise KeyError(f"missing inputs for {self.path.name}: {missing}")
        prepared = {
            name: np.ascontiguousarray(np.asarray(feed[name], dtype=meta.dtype))
            for name, meta in self.input_by_name.items()
        }
        values = self._session.run(None, prepared)
        if isinstance(values, dict):
            return {name: np.asarray(value) for name, value in values.items()}
        if not isinstance(values, (list, tuple)):
            values = [values]
        if len(values) != len(self.outputs):
            raise RuntimeError(
                f"unexpected output count from {self.path.name}: "
                f"{len(values)} != {len(self.outputs)}"
            )
        return {
            meta.name: np.asarray(value)
            for meta, value in zip(self.outputs, values)
        }


def first_output(session: InferenceSession, outputs: dict[str, np.ndarray]) -> np.ndarray:
    return outputs[session.outputs[0].name]