#!/usr/bin/env python3

from __future__ import annotations

import logging
from pathlib import Path
from typing import Dict, List

import numpy as np
import onnxruntime as ort

from scripts.zipvoice_decoder4_runtime import Decoder4ZipVoiceBoardRuntime
from scripts.zipvoice_runtime import AxeSession


class _OrtSession:
    """ONNX Runtime wrapper with the AxeSession interface used by split runtime."""

    _TYPE_TO_DTYPE = {
        "tensor(float)": np.float32,
        "tensor(float32)": np.float32,
        "tensor(double)": np.float64,
        "tensor(int32)": np.int32,
        "tensor(int64)": np.int64,
        "tensor(uint8)": np.uint8,
        "tensor(bool)": np.bool_,
    }

    def __init__(self, model_path: str | Path):
        self.path = Path(model_path)
        if not self.path.exists():
            raise FileNotFoundError(f"ONNX model not found: {self.path}")
        self._session = ort.InferenceSession(
            str(self.path),
            providers=["CPUExecutionProvider"],
        )
        self._inputs = self._session.get_inputs()
        self._outputs = self._session.get_outputs()

    @property
    def input_names(self) -> List[str]:
        return [item.name for item in self._inputs]

    @property
    def output_names(self) -> List[str]:
        return [item.name for item in self._outputs]

    def _coerce_one(self, name: str, value: np.ndarray) -> np.ndarray:
        info = next((item for item in self._inputs if item.name == name), None)
        array = np.asarray(value)
        if info is None:
            return np.ascontiguousarray(array)

        dtype = self._TYPE_TO_DTYPE.get(str(info.type).lower())
        if dtype is not None:
            array = array.astype(dtype, copy=False)

        if list(info.shape) == [] and array.shape == (1,):
            array = array.reshape(())

        return np.ascontiguousarray(array)

    def run(self, feed_dict: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        feed = {name: self._coerce_one(name, value) for name, value in feed_dict.items()}
        outputs = self._session.run(None, feed)
        return {name: value for name, value in zip(self.output_names, outputs)}


class Decoder4ZipVoiceBoardRuntimeEncoderOnnx(Decoder4ZipVoiceBoardRuntime):
    """Runs encoder with ONNX Runtime, all decoder parts with axmodel."""

    def _load_models(self) -> None:
        self.sessions = {}

        encoder_name = self.encoder_info["name"]
        encoder_path = self.models_dir / "encoder_core.onnx"
        logging.info("encoder 使用 ONNX Runtime: %s", encoder_path)
        self.sessions[encoder_name] = _OrtSession(encoder_path)

        for info in self.decoder_parts:
            name = info["name"]
            path = self.models_dir / info["file"]
            logging.debug("Loading %s from %s", name, path)
            self.sessions[name] = AxeSession(path)

        self.decoder_label = "decoder4(encoder_onnx+part0-3_axmodel)"
        logging.debug("Loaded encoder ONNX + %d decoder axmodels", len(self.decoder_parts))