"""3 个 axmodel 的 axengine 会话封装 (enc/erb_dec/df_dec 流水线)。"""
from __future__ import annotations

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

import numpy as np

T = 99  # 静态时间帧 (1s @ 48kHz 流式 STFT)


@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 _AxSession:
    """单个 axmodel 的 axengine 会话 (参考 AXERA 官方推理工程封装)。"""

    def __init__(self, model_path: str | Path):
        import axengine
        self.path = Path(model_path)
        if not self.path.is_file():
            raise FileNotFoundError(self.path)
        self._session = axengine.InferenceSession(str(self.path))
        self.inputs = [_tensor_info(v) for v in self._session.get_inputs()]
        self.outputs = [_tensor_info(v) for v in self._session.get_outputs()]
        self.input_by_name = {v.name: v for v in self.inputs}

    def run(self, feed: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
        missing = [v.name for v in self.inputs if v.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(v) for name, v 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(v) for meta, v in zip(self.outputs, values)}


class DfNet3Ax:
    """DeepFilterNet3 三子图 NPU 流水线 (与 dsp.pipeline 配合)。

    enc:    feat_erb(1,1,T,32)+feat_spec(1,2,T,96) -> e0..e3, emb, c0, lsnr
    erb_dec: emb,e3,e2,e1,e0 -> m(1,1,T,32)
    df_dec:  emb,c0 -> coefs(1,T,96,10)
    """

    def __init__(self, model_dir: str | Path):
        model_dir = Path(model_dir)
        self.enc = _AxSession(model_dir / "enc.axmodel")
        self.erb_dec = _AxSession(model_dir / "erb_dec.axmodel")
        self.df_dec = _AxSession(model_dir / "df_dec.axmodel")
        # 状态携带模式检测: axmodel 是否有 h0 输入
        self.stateful = all("h0" in s.input_by_name for s in (self.enc, self.df_dec))

    def run_enc(self, feat_erb: np.ndarray, feat_spec: np.ndarray,
                h: np.ndarray | None = None) -> tuple[dict[str, np.ndarray], np.ndarray | None]:
        feed = {"feat_erb": feat_erb, "feat_spec": feat_spec}
        if self.stateful:
            if h is None:
                h = np.zeros((1, 1, 256), np.float32)
            feed["h0"] = h
            out = self.enc.run(feed)
            return out, np.asarray(out["h_out"])
        return self.enc.run(feed), None

    def run_erb_dec(self, emb, e3, e2, e1, e0) -> np.ndarray:
        return self.erb_dec.run({"emb": emb, "e3": e3, "e2": e2, "e1": e1, "e0": e0})["m"]

    def run_df_dec(self, emb, c0, h: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray | None]:
        feed = {"emb": emb, "c0": c0}
        if self.stateful:
            if h is None:
                h = np.zeros((2, 1, 256), np.float32)
            feed["h0"] = h
            out = self.df_dec.run(feed)
            return np.asarray(out["coefs"]), np.asarray(out["h_out"])
        return self.df_dec.run(feed)["coefs"], None