"""
VibeVoice ASR engine v4 — Pure VibeVoice source (model + processor).
Patched AutoModel.register(..., exist_ok=True) for Transformers 5.x compat.
"""
import torch
import sys
sys.path.insert(0, "/share/ymq/VibeVoice")

from longtasks.longtasks import LongTasks
from appPublic.worker import awaitify
from appPublic.jsonConfig import getConfig
from appPublic.log import debug
from ahserver.filestorage import FileStorage

from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor


class VibeVoiceASREngine(LongTasks):
    def __init__(self):
        self.config = getConfig()
        super().__init__(self.config.redis_url, 'vibevoice-asr', worker_cnt=self.config.worker_cnt)
        self.load_models()

    def load_models(self):
        model_path = self.config.model_path
        debug('loading VibeVoice ASR on cuda:0...')
        self.model = VibeVoiceASRForConditionalGeneration.from_pretrained(
            model_path,
            dtype=torch.bfloat16,
            device_map="cuda:0",
            attn_implementation="sdpa",
            trust_remote_code=True,
            ignore_mismatched_sizes=True,
        )
        self.model.eval()
        self.processor = VibeVoiceASRProcessor.from_pretrained(
            model_path,
            language_model_pretrained_name="/share/models/Qwen2.5-7B",
        )
        debug('VibeVoice ASR loaded')

    async def process_task(self, payload, workerid=None):
        webpath = payload.get('audio_file')
        if not webpath:
            return {"task_status": "error", "message": "no audio_file"}
        fs = FileStorage()
        fpath = fs.realPath(webpath)
        f = awaitify(self._transcribe)
        return await f(fpath)

    def _transcribe(self, fpath):
        inputs = self.processor(
            audio=fpath, sampling_rate=None, return_tensors="pt",
            padding=True, add_generation_prompt=True)
        device = torch.device("cuda:0")
        inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v
                  for k, v in inputs.items()}
        with torch.no_grad():
            gen = self.model.generate(**inputs, max_new_tokens=4096,
                                       temperature=0.0, do_sample=False)
        text = self.processor.decode(gen[0], skip_special_tokens=True)
        try:
            segs = self.processor.post_process_transcription(text)
        except Exception as e:
            debug(f'parse failed: {e}')
            segs = []

        # Fallback: direct JSON extraction from model output
        if not segs and text:
            import re, json as _json
            # Try full JSON first
            m = re.search(r'assistant\n(.+?)(?:<\|im_end\|>)', text, re.DOTALL)
            json_str = m.group(1).strip() if m else ''
            if not json_str and 'assistant\n[' in text:
                json_str = text[text.find('assistant\n[')+10:]
            if json_str:
                # Handle truncated JSON: find last complete object
                try:
                    raw_segs = _json.loads(json_str)
                except Exception:
                    last_good = json_str.rfind('"}')
                    if last_good > 0:
                        try:
                            raw_segs = _json.loads(json_str[:last_good+2] + ']')
                        except Exception:
                            raw_segs = []
                    else:
                        raw_segs = []
                for s in raw_segs:
                    segs.append({
                        'start_time': s.get('Start', 0),
                        'end_time': s.get('End', 0),
                        'speaker_id': s.get('Speaker', 0),
                        'text': s.get('Content', ''),
                    })

        ws, ft = [], []
        for s in segs:
            ws.append([s.get('start_time',0), s.get('end_time',0), s.get('text',''), []])
            ft.append(s.get('text',''))
        return {
            'task_status': 'SUCCEEDED', 'language': 'auto',
            'language_probability': 1.0, 'content': ' '.join(ft),
            'segments': ws, 'raw_output': text, 'vibevoice_segments': segs,
        }