import sys, os, re, torch, torchaudio, numpy as np
from typing import Dict, Any


class EndpointHandler:
    def __init__(self, path=""):
        if path and path not in sys.path:
            sys.path.insert(0, path)

        from configuration_moss_audio import MossAudioConfig
        from processing_moss_audio import MossAudioProcessor
        from modeling_moss_audio import MossAudioModel

        print(f"[MOSS-8B] Loading (no time markers)...")
        self.config = MossAudioConfig.from_pretrained(path)
        self.processor = MossAudioProcessor.from_pretrained(
            path, trust_remote_code=True, enable_time_marker=False
        )
        self.model = MossAudioModel.from_pretrained(
            path, config=self.config, torch_dtype=torch.bfloat16,
            device_map="auto", trust_remote_code=True,
        )
        self.model.eval()
        print(f"[MOSS-8B] Loaded.")
        self.sr = self.processor.config.mel_sr

    def _load_audio(self, audio_input):
        if isinstance(audio_input, str) and audio_input.startswith(("http://", "https://")):
            import requests, tempfile
            resp = requests.get(audio_input, timeout=120)
            resp.raise_for_status()
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
                tmp.write(resp.content)
                tmp_path = tmp.name
            waveform, sr = torchaudio.load(tmp_path)
            os.unlink(tmp_path)
        elif isinstance(audio_input, str):
            waveform, sr = torchaudio.load(audio_input)
        else:
            raise ValueError("audio must be URL or file path")

        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
        if sr != self.sr:
            waveform = torchaudio.transforms.Resample(sr, self.sr)(waveform)
        return waveform

    def _transcribe(self, raw_audio, prompt, gen_kwargs):
        inputs = self.processor(text=prompt, audios=[raw_audio], return_tensors="pt")
        inputs = inputs.to(self.model.device)
        if inputs.get("audio_data") is not None:
            inputs["audio_data"] = inputs["audio_data"].to(self.model.dtype)
        inputs["audio_input_mask"] = inputs["input_ids"] == self.processor.audio_token_id
        ids = self.model.generate(**inputs, **gen_kwargs)
        text = self.processor.decode(ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
        text = re.sub(r'<think>.*?</think>\s*\n?\s*', '', text, flags=re.DOTALL).strip()
        return text

    @torch.no_grad()
    def __call__(self, data):
        inputs = data.get("inputs", data)
        try:
            if isinstance(inputs, str):
                audio_input, prompt = inputs, "转写这段日语音频，标注说话人。格式：说话人：内容"
            elif isinstance(inputs, dict):
                audio_input = inputs.get("audio", "")
                prompt = inputs.get("prompt", inputs.get("text", "转写这段日语音频，标注说话人。格式：说话人：内容"))
            else:
                return {"error": f"Invalid input type: {type(inputs)}"}

            gen_kwargs = dict(
                max_new_tokens=int(data.get("max_new_tokens", inputs.get("max_new_tokens", 384))),
                num_beams=int(data.get("num_beams", inputs.get("num_beams", 2))),
                use_cache=True,
                do_sample=bool(data.get("do_sample", inputs.get("do_sample", False))),
                temperature=float(data.get("temperature", inputs.get("temperature", 0.5))),
                top_p=float(data.get("top_p", inputs.get("top_p", 0.9))),
                top_k=int(data.get("top_k", inputs.get("top_k", 50))),
                repetition_penalty=float(data.get("repetition_penalty", inputs.get("repetition_penalty", 1.15))),
                no_repeat_ngram_size=int(data.get("no_repeat_ngram_size", inputs.get("no_repeat_ngram_size", 3))),
                length_penalty=float(data.get("length_penalty", inputs.get("length_penalty", 0.9))),
                early_stopping=True,
            )

            waveform = self._load_audio(audio_input)
            total_sec = waveform.shape[1] / self.sr
            print(f"[MOSS-8B] Audio: {total_sec:.1f}s (no time markers)")

            # 2-minute (120s) chunks with 10s overlap
            chunk_sec = 120
            overlap_sec = 10
            step_samples = int((chunk_sec - overlap_sec) * self.sr)
            chunk_samples = int(chunk_sec * self.sr)
            n_frames = waveform.shape[1]

            chunks_list = []
            pos = 0
            while pos < n_frames:
                end = min(pos + chunk_samples, n_frames)
                chunks_list.append((pos, end))
                pos += step_samples
                if end >= n_frames:
                    break

            segments = []
            prev_speakers = ""

            for idx, (st, en) in enumerate(chunks_list):
                chunk = waveform[:, st:en]
                if chunk.abs().max() < 0.005:
                    text = ""
                else:
                    ctx_prompt = prompt
                    if prev_speakers and idx > 0:
                        ctx_prompt = prompt + f" 前段说话人：{prev_speakers}"

                    raw = chunk.squeeze(0).numpy().astype(np.float32)
                    text = self._transcribe(raw, ctx_prompt, gen_kwargs)

                    speakers = set(re.findall(r'([\u3040-\u30ff\u4e00-\u9fff]+)[：:]', text))
                    if speakers:
                        prev_speakers = "、".join(sorted(speakers)[:4])

                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()

                start_t = st / self.sr
                end_t = en / self.sr
                segments.append({"start": round(start_t, 2), "end": round(end_t, 2), "text": text})
                print(f"[Chunk {idx+1}/{len(chunks_list)}] ({len(text)}c) {text[:80]}")

            full_text = " ".join(s["text"] for s in segments if s["text"])

            srt_lines = []
            idx_n = 1
            for seg in segments:
                t = seg["text"].strip()
                if not t:
                    continue
                s, e = seg["start"], seg["end"]
                srt_lines.append(str(idx_n))
                srt_lines.append(f"{int(s//3600):02d}:{int(s%3600//60):02d}:{s%60:06.3f} --> {int(e//3600):02d}:{int(e%3600//60):02d}:{e%60:06.3f}")
                srt_lines.append(t)
                srt_lines.append("")
                idx_n += 1

            return {
                "transcription": full_text,
                "segments": segments,
                "srt": "\n".join(srt_lines),
                "num_chunks": len(chunks_list),
            }

        except Exception as e:
            import traceback
            return {"error": str(e), "traceback": traceback.format_exc()}