#!/usr/bin/env python3
from __future__ import annotations

import json
import wave
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(int16)", "int16", "s16"), np.int16),
        (("tensor(int8)", "int8", "s8"), np.int8),
        (("tensor(uint16)", "uint16", "u16"), np.uint16),
        (("tensor(uint8)", "uint8", "u8"), np.uint8),
        (("tensor(bool)", "bool"), np.bool_),
    )
    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)
    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:
    """Small common wrapper for ONNX Runtime and AXEngine."""

    def __init__(self, model_path: str | Path, backend: str):
        self.path = Path(model_path)
        if not self.path.is_file():
            raise FileNotFoundError(self.path)
        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]


def read_wav(path: str | Path) -> tuple[np.ndarray, int]:
    path = Path(path)
    with wave.open(str(path), "rb") as source:
        if source.getnchannels() != 1 or source.getsampwidth() != 2:
            raise ValueError(f"expected mono 16-bit PCM WAV: {path}")
        sample_rate = source.getframerate()
        samples = np.frombuffer(
            source.readframes(source.getnframes()), dtype=np.int16
        ).copy()
    return samples, sample_rate


def load_json(path: str | Path) -> dict[str, Any]:
    return json.loads(Path(path).read_text())


def write_json(path: str | Path, value: Any) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n")


def max_abs_delta(reference: Any, candidate: Any) -> float:
    reference_array = np.asarray(reference, dtype=np.float64)
    candidate_array = np.asarray(candidate, dtype=np.float64)
    if reference_array.shape != candidate_array.shape:
        return float("inf")
    if reference_array.size == 0:
        return 0.0
    return float(np.max(np.abs(reference_array - candidate_array)))