# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Triton Python backend: the release's `load_audio` resampling (torchaudio 2.1.1's sinc resampler, as copied unchanged
into anuvani_diar.nn.resample, 2 CPU threads) for non-16 kHz input."""
import numpy as np
import triton_python_backend_utils as pb_utils


class TritonPythonModel:
    def initialize(self, args):
        import torch
        from anuvani_diar.nn.resample import resample
        self.torch, self.resample = torch, resample
        torch.set_num_threads(2)

    def execute(self, requests):
        t, resample = self.torch, self.resample
        responses = []
        for req in requests:
            try:
                audio = pb_utils.get_input_tensor_by_name(req, "AUDIO").as_numpy()
                sr = int(pb_utils.get_input_tensor_by_name(req, "SR").as_numpy()[0])
                t.set_num_threads(2)
                out = resample(t.from_numpy(np.ascontiguousarray(audio)).unsqueeze(0), sr, 16000).squeeze(0).numpy().astype(np.float32)
                responses.append(pb_utils.InferenceResponse(output_tensors=[pb_utils.Tensor("OUT", out)]))
            except Exception as e:
                responses.append(pb_utils.InferenceResponse(output_tensors=[], error=pb_utils.TritonError(f"{type(e).__name__}: {e}")))
        return responses