"""ECAPA-TDNN Voiceprint Engine — loads model, processes extract/verify tasks."""
import torch
from longtasks.longtasks import LongTasks
from appPublic.worker import awaitify
from appPublic.jsonConfig import getConfig
from appPublic.log import debug
from speechbrain.inference.speaker import SpeakerRecognition
import numpy as np


class VoiceprintEngine(LongTasks):
    def __init__(self):
        self.config = getConfig()
        super().__init__(self.config.redis_url, 'voiceprint', worker_cnt=self.config.worker_cnt)
        self.load_model()

    def load_model(self):
        device = self.config.device
        debug(f'loading ECAPA-TDNN on {device}...')
        self.verifier = SpeakerRecognition.from_hparams(
            source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="/share/models/ecapa-tdnn",
            run_opts={"device": device}
        )
        debug('ECAPA-TDNN loaded')

    async def process_task(self, payload, workerid=None):
        task_type = payload.get('task_type', 'extract')
        audio_file = payload.get('audio_file', '')
        if not audio_file:
            return {'status': 'FAILED', 'result': 'missing audio_file'}

        if task_type == 'extract':
            f = awaitify(self._extract)
            return await f(audio_file)
        elif task_type == 'verify':
            ref_file = payload.get('reference_file', '')
            if not ref_file:
                return {'status': 'FAILED', 'result': 'missing reference_file'}
            f = awaitify(self._verify)
            return await f(audio_file, ref_file)
        return {'status': 'FAILED', 'result': f'unknown task_type: {task_type}'}

    def _extract(self, audio_path):
        signal = self.verifier.load_audio(audio_path, 16000)
        t = torch.tensor(signal).unsqueeze(0).to(self.config.device)
        emb = self.verifier.encode_batch(t)
        vec = emb.squeeze().cpu().numpy().tolist()
        return {
            'status': 'SUCCEEDED',
            'embedding': vec,
            'embedding_dim': len(vec),
            'usage': {'audio_duration': round(len(signal) / 16000, 2)},
        }

    def _verify(self, audio_path, ref_path):
        score, pred = self.verifier.verify_files(audio_path, ref_path)
        return {
            'status': 'SUCCEEDED',
            'similarity': round(float(score), 4),
            'is_same_speaker': bool(pred),
            'usage': {},
        }