import json
from pathlib import Path
from typing import Any

import torch
from torch import nn

from .configuration_onnx_parameter_store import OnnxParameterStoreConfig


class OnnxParameterStoreModel(nn.Module):
    config_class = OnnxParameterStoreConfig
    base_model_prefix = "onnx_parameter_store"
    _auto_class = "AutoModel"

    def __init__(self, config: OnnxParameterStoreConfig):
        super().__init__()
        self.config = config
        self.weights = torch.nn.ParameterDict()
        self._onnx_session = None
        self._onnx_path = None

        spec_path = self._resolve_parameter_spec_path()
        with spec_path.open("r", encoding="utf-8") as f:
            items = json.load(f)["items"]

        for item in items:
            name = item["pt_name"]
            shape = item["shape"]
            dtype = _str_to_dtype(item["dtype"])
            tensor = torch.zeros(shape, dtype=dtype)
            self.weights[name] = torch.nn.Parameter(tensor, requires_grad=False)

    def forward(
        self,
        parameter_name: str | None = None,
        audio: Any | None = None,
        audio_path: str | None = None,
        sample_rate: int | None = None,
        onnx_path: str | None = None,
    ):
        if parameter_name is not None:
            return {"parameter": self.weights[parameter_name]}

        if audio is not None:
            indices = self.tokenize(audio=audio, sample_rate=sample_rate, onnx_path=onnx_path)
            return {"indices": indices}

        if audio_path is not None:
            indices = self.tokenize_from_file(audio_path=audio_path, onnx_path=onnx_path, sample_rate=sample_rate)
            return {"indices": indices}

        return {
            "num_parameters": torch.tensor(len(self.weights), dtype=torch.long),
            "keys": list(self.weights.keys()),
        }

    def tokenize_from_file(
        self,
        audio_path: str,
        onnx_path: str | None = None,
        sample_rate: int | None = None,
    ) -> torch.Tensor:
        try:
            import soundfile as sf
        except Exception as exc:
            raise RuntimeError("soundfile가 필요합니다. `pip install soundfile` 후 재시도하세요.") from exc

        wav, sr = sf.read(str(audio_path), dtype="float32")
        if wav.ndim == 2:
            wav = wav.mean(axis=1)
        src_sr = int(sample_rate) if sample_rate is not None else int(sr)
        return self.tokenize(audio=wav, sample_rate=src_sr, onnx_path=onnx_path)

    def tokenize(
        self,
        audio: Any,
        sample_rate: int | None = None,
        onnx_path: str | None = None,
    ) -> torch.Tensor:
        try:
            import numpy as np
        except Exception as exc:
            raise RuntimeError("numpy가 필요합니다. `pip install numpy` 후 재시도하세요.") from exc

        if isinstance(audio, torch.Tensor):
            wav = audio.detach().cpu().numpy()
        else:
            wav = np.asarray(audio)

        wav = np.asarray(wav, dtype=np.float32)
        if wav.ndim == 2:
            wav = wav.mean(axis=1)

        src_sr = int(sample_rate) if sample_rate is not None else int(self.config.audio_sample_rate)
        tgt_sr = int(self.config.audio_sample_rate)
        if src_sr != tgt_sr:
            wav = self._resample_linear(wav, src_sr=src_sr, tgt_sr=tgt_sr)

        feats = self._extract_logmel(wav)
        feats_len = np.asarray([feats.shape[-1]], dtype=np.int32)

        session = self._get_onnx_session(onnx_path=onnx_path)
        input_names = [i.name for i in session.get_inputs()]
        feats_name = next((n for n in input_names if "feat" in n.lower()), input_names[0])
        len_name = next((n for n in input_names if "length" in n.lower()), input_names[-1])
        outputs = session.run(None, {feats_name: feats.astype(np.float32), len_name: feats_len})
        indices = np.asarray(outputs[0], dtype=np.int64)
        return torch.from_numpy(indices)

    def _get_onnx_session(self, onnx_path: str | None = None):
        try:
            import onnxruntime as ort
        except Exception as exc:
            raise RuntimeError("onnxruntime가 필요합니다. `pip install onnxruntime` 후 재시도하세요.") from exc

        resolved = str(Path(onnx_path).expanduser().resolve()) if onnx_path else self._resolve_onnx_path()
        if self._onnx_session is not None and self._onnx_path == resolved:
            return self._onnx_session

        self._onnx_session = ort.InferenceSession(resolved, providers=["CPUExecutionProvider"])
        self._onnx_path = resolved
        return self._onnx_session

    def _resolve_parameter_spec_path(self) -> Path:
        candidates = []
        name_or_path = str(getattr(self.config, "_name_or_path", "")).strip()
        repo_root = Path(name_or_path)
        if name_or_path and repo_root.is_dir():
            candidates.append(repo_root / self.config.parameter_spec_file)
        candidates.append(Path(__file__).resolve().parent / self.config.parameter_spec_file)

        for cand in candidates:
            if cand.is_file():
                return cand

        if name_or_path and "/" in name_or_path:
            try:
                from huggingface_hub import hf_hub_download
                downloaded = hf_hub_download(repo_id=name_or_path, filename=self.config.parameter_spec_file)
                cand = Path(downloaded)
                if cand.is_file():
                    return cand
            except Exception:
                pass

        raise FileNotFoundError(
            f"Cannot find parameter spec file: {self.config.parameter_spec_file}. candidates={candidates}"
        )

    def _resolve_onnx_path(self) -> str:
        candidates = []
        local_repo = str(getattr(self.config, "_name_or_path", "")).strip()
        if local_repo:
            local_repo_path = Path(local_repo)
            if local_repo_path.is_dir():
                candidates.append(local_repo_path / self.config.source_onnx_file)
        candidates.append(Path(__file__).resolve().parent / self.config.source_onnx_file)

        for cand in candidates:
            if cand.is_file():
                return str(cand.resolve())

        repo_id = self.config.source_onnx_repo_id
        if not repo_id and local_repo and _looks_like_repo_id(local_repo):
            repo_id = local_repo
        if repo_id and "/" in repo_id:
            try:
                from huggingface_hub import hf_hub_download
                downloaded = hf_hub_download(repo_id=repo_id, filename=self.config.source_onnx_file)
                return str(Path(downloaded).resolve())
            except Exception as exc:
                raise FileNotFoundError(
                    "Cannot resolve ONNX file. "
                    f"repo_id={repo_id}, filename={self.config.source_onnx_file}, candidates={candidates}"
                ) from exc

        raise FileNotFoundError(
            "Cannot resolve ONNX file. Provide `onnx_path` when calling tokenize(), "
            f"or set config.source_onnx_repo_id/source_onnx_file. candidates={candidates}"
        )

    def _extract_logmel(self, wav):
        import numpy as np

        n_fft = int(self.config.n_fft)
        win_length = int(self.config.win_length)
        hop_length = int(self.config.hop_length)
        n_mels = int(self.config.n_mels)
        sr = int(self.config.audio_sample_rate)

        if wav.shape[0] < win_length:
            wav = np.pad(wav, (0, win_length - wav.shape[0]))

        frames = 1 + (wav.shape[0] - win_length) // hop_length
        if frames <= 0:
            frames = 1
        pad_len = win_length + hop_length * (frames - 1)
        if pad_len > wav.shape[0]:
            wav = np.pad(wav, (0, pad_len - wav.shape[0]))

        idx = np.arange(win_length)[None, :] + hop_length * np.arange(frames)[:, None]
        framed = wav[idx]
        window = np.hanning(win_length).astype(np.float32)
        spectrum = np.fft.rfft(framed * window[None, :], n=n_fft, axis=1)
        power = (np.abs(spectrum) ** 2).astype(np.float32)
        fb = self._mel_filterbank(sr=sr, n_fft=n_fft, n_mels=n_mels)
        mel = power @ fb.T
        logmel = np.log(np.clip(mel, 1e-10, None)).astype(np.float32)
        return logmel.T[None, :, :]

    def _resample_linear(self, x, src_sr: int, tgt_sr: int):
        import numpy as np
        if src_sr == tgt_sr or x.size < 2:
            return x.astype(np.float32, copy=False)
        new_len = max(1, int(round(x.shape[0] * float(tgt_sr) / float(src_sr))))
        old = np.arange(x.shape[0], dtype=np.float64)
        new = np.linspace(0, x.shape[0] - 1, num=new_len, dtype=np.float64)
        return np.interp(new, old, x.astype(np.float64)).astype(np.float32)

    def _mel_filterbank(self, sr: int, n_fft: int, n_mels: int, f_min: float = 0.0, f_max: float | None = None):
        import numpy as np
        if f_max is None:
            f_max = sr / 2.0
        mel_min = self._hz_to_mel(np.array([f_min]))[0]
        mel_max = self._hz_to_mel(np.array([f_max]))[0]
        mels = np.linspace(mel_min, mel_max, n_mels + 2)
        hz = self._mel_to_hz(mels)
        bins = np.floor((n_fft + 1) * hz / sr).astype(int)
        fb = np.zeros((n_mels, n_fft // 2 + 1), dtype=np.float32)
        for i in range(1, n_mels + 1):
            left, center, right = bins[i - 1], bins[i], bins[i + 1]
            if center > left:
                fb[i - 1, left:center] = (np.arange(left, center) - left) / float(center - left)
            if right > center:
                fb[i - 1, center:right] = (right - np.arange(center, right)) / float(right - center)
        return fb

    def _hz_to_mel(self, hz):
        import numpy as np
        return 2595.0 * np.log10(1.0 + hz / 700.0)

    def _mel_to_hz(self, mel):
        return 700.0 * (10.0 ** (mel / 2595.0) - 1.0)

    @classmethod
    def register_for_auto_class(cls, auto_class: str = "AutoModel"):
        cls._auto_class = auto_class

    @property
    def device(self) -> torch.device:
        try:
            return next(self.parameters()).device
        except StopIteration:
            return torch.device("cpu")

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | Path,
        *model_args,
        config: OnnxParameterStoreConfig | None = None,
        **kwargs,
    ):
        if model_args:
            raise ValueError("OnnxParameterStoreModel does not use positional model args.")

        resolved_dir = cls._resolve_repo_dir(
            pretrained_model_name_or_path=pretrained_model_name_or_path,
            revision=kwargs.get("revision"),
            cache_dir=kwargs.get("cache_dir"),
            token=kwargs.get("token"),
            local_files_only=bool(kwargs.get("local_files_only", False)),
        )

        if config is None:
            try:
                from transformers import AutoConfig
            except Exception as exc:
                raise RuntimeError("transformers가 필요합니다. `pip install transformers` 후 재시도하세요.") from exc
            config = AutoConfig.from_pretrained(
                str(resolved_dir),
                trust_remote_code=True,
                revision=kwargs.get("revision"),
                cache_dir=kwargs.get("cache_dir"),
                token=kwargs.get("token"),
                local_files_only=bool(kwargs.get("local_files_only", False)),
            )
        setattr(config, "_name_or_path", str(resolved_dir))

        model = cls(config)
        state_path = cls._resolve_state_dict_path(resolved_dir)
        state_dict = cls._load_state_dict(state_path)
        model.load_state_dict(state_dict, strict=True)
        model.eval()
        return model

    @staticmethod
    def _resolve_repo_dir(
        pretrained_model_name_or_path: str | Path,
        revision: str | None = None,
        cache_dir: str | None = None,
        token: str | None = None,
        local_files_only: bool = False,
    ) -> Path:
        local = Path(str(pretrained_model_name_or_path)).expanduser()
        if local.is_dir():
            return local.resolve()

        repo_id = str(pretrained_model_name_or_path)
        if "/" not in repo_id:
            raise FileNotFoundError(
                f"Cannot find local directory: {local}. "
                "For remote loading, pass a valid Hugging Face repo id like `user/model`."
            )

        try:
            from huggingface_hub import snapshot_download
        except Exception as exc:
            raise RuntimeError(
                "huggingface_hub가 필요합니다. `pip install huggingface_hub` 후 재시도하세요."
            ) from exc

        downloaded = snapshot_download(
            repo_id=repo_id,
            revision=revision,
            cache_dir=cache_dir,
            token=token,
            local_files_only=local_files_only,
        )
        return Path(downloaded).resolve()

    @staticmethod
    def _resolve_state_dict_path(repo_dir: Path) -> Path:
        for name in ("pytorch_model.bin", "model.safetensors"):
            cand = repo_dir / name
            if cand.is_file():
                return cand
        raise FileNotFoundError(f"Cannot find model weights in {repo_dir}.")

    @staticmethod
    def _load_state_dict(path: Path) -> dict[str, torch.Tensor]:
        if path.suffix == ".safetensors":
            try:
                from safetensors.torch import load_file
            except Exception as exc:
                raise RuntimeError(
                    "safetensors 로딩을 위해 `pip install safetensors`가 필요합니다."
                ) from exc
            return load_file(str(path))
        return torch.load(path, map_location="cpu")


def _str_to_dtype(name: str):
    table = {
        "float32": torch.float32,
        "float16": torch.float16,
        "float64": torch.float64,
        "int64": torch.int64,
        "int32": torch.int32,
        "int16": torch.int16,
        "int8": torch.int8,
        "uint8": torch.uint8,
        "bool": torch.bool,
    }
    return table.get(name, torch.float32)


def _looks_like_repo_id(candidate: str) -> bool:
    path_like = Path(candidate)
    if path_like.exists():
        return False
    if candidate.startswith("/") or candidate.startswith("."):
        return False
    if "\\" in candidate or ":" in candidate:
        return False
    return "/" in candidate