"""DHVANI proprietary custom model runner."""

from __future__ import annotations

import json
import logging
from pathlib import Path

import numpy as np
import torch

from core.config import SAMPLE_RATE, custom_model_path
from core.dhvani_head import BACKBONE_ID, DhvaniCustomModel, MAX_SAMPLES

logger = logging.getLogger("dhvani.custom")


class CustomDhvaniDetector:
    """Loads models/dhvani-custom.pt trained via training/train_dhvani_head.py."""

    def __init__(self, checkpoint_path: str | None = None, device: str | None = None) -> None:
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.checkpoint_path = checkpoint_path or custom_model_path()
        self._model: DhvaniCustomModel | None = None
        self._backbone_id = BACKBONE_ID
        self._fake_label_index = 1
        self._metadata: dict = {}

    @property
    def enabled(self) -> bool:
        return bool(self.checkpoint_path) and Path(self.checkpoint_path).exists()

    def load(self) -> None:
        if not self.enabled:
            return
        if self._model is not None:
            return

        logger.info("Loading DHVANI custom model from %s", self.checkpoint_path)
        payload = torch.load(self.checkpoint_path, map_location="cpu", weights_only=False)
        self._backbone_id = payload.get("backbone_id", self._backbone_id)
        self._fake_label_index = int(payload.get("fake_label_index", 1))
        self._metadata = {k: v for k, v in payload.items() if k != "head_state_dict"}

        self._model = DhvaniCustomModel(backbone_id=self._backbone_id)
        self._model.head.load_state_dict(payload["head_state_dict"])
        self._model.to(self.device)
        self._model.eval()

    @property
    def loaded(self) -> bool:
        return self._model is not None

    @property
    def model_id(self) -> str:
        version = self._metadata.get("version", "")
        if version:
            return str(version)
        return "dhvani-custom-head"

    @property
    def version(self) -> str:
        return str(self._metadata.get("version", "unknown"))

    @property
    def metadata(self) -> dict:
        return dict(self._metadata)

    def fake_probability(self, waveform: np.ndarray) -> float:
        if not self.enabled:
            return 0.0
        self.load()
        assert self._model is not None

        if waveform.size > MAX_SAMPLES:
            waveform = waveform[:MAX_SAMPLES]
        if waveform.size < SAMPLE_RATE:
            waveform = np.pad(waveform, (0, SAMPLE_RATE - waveform.size))

        tensor = torch.from_numpy(waveform.astype(np.float32)).unsqueeze(0).to(self.device)
        with torch.no_grad():
            logits = self._model(tensor)
            probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
        return float(probs[self._fake_label_index])