import numpy as np
import noisereduce as nr
import torch
import io
import time
import wave
import threading
import logging

from groq import Groq
from optimum.onnxruntime import ORTModelForSpeechSeq2Seq
from transformers import WhisperProcessor

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

SAMPLE_RATE     = 16000
GROQ_TIMEOUT    = 5.0
VAD_THRESHOLD   = 0.3
BUFFER_DURATION = 2.5
BUFFER_SAMPLES  = int(SAMPLE_RATE * BUFFER_DURATION)

# Confidence threshold — kalau avg_logprob < nilai ini, fallback ke Groq
CONFIDENCE_THRESHOLD = -0.6

HF_MODEL_ID = "Isaacyn/whisper-small-id"

GROQ_API_KEYS = [
    "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",
    # tambah key lain di sini kalau ada
]
_groq_key_idx = 0
_groq_key_lock = threading.Lock()

def _get_groq_client():
    global _groq_key_idx
    with _groq_key_lock:
        key = GROQ_API_KEYS[_groq_key_idx % len(GROQ_API_KEYS)]
        _groq_key_idx += 1
    return Groq(api_key=key)

BLACKLIST = [
    "transkripsi percakapan", "bahasa indonesia", "transkripsi",
    "indonesia.", "indonesia", "terima kasih", "thank you",
    "thanks", "subscribe", "like", "comment",
]

EXACT_BLACKLIST = [
    "indonesia", "terima kasih", "thank you", "thanks",
    "like", "subscribe", "comment", "ok", "oke",
]

ENGLISH_STARTERS = [
    "so ", "we ", "the ", "this ", "that ", "i ", "you ", "it ",
    "and ", "but ", "for ", "are ", "is ", "was ", "in ", "of ",
    "a ", "an ", "to ", "with ", "my ", "your ", "our ", "they ",
    "he ", "she ", "what ", "how ", "when ", "where ", "why ",
]


def audio_to_wav_bytes(audio: np.ndarray, sample_rate: int = SAMPLE_RATE) -> bytes:
    buf = io.BytesIO()
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        pcm = (audio * 32767).astype(np.int16)
        wf.writeframes(pcm.tobytes())
    return buf.getvalue()


class ASRPipeline:
    def __init__(self, model_dir: str = None):
        # ── Whisper ONNX (primary) ──────────────────────────────────────
        logger.info(f"Loading Whisper ONNX from {HF_MODEL_ID} ...")
        self.processor = WhisperProcessor.from_pretrained(HF_MODEL_ID)
        self.onnx_model = ORTModelForSpeechSeq2Seq.from_pretrained(
            HF_MODEL_ID,
            export=False,           # sudah ada ONNX di repo
            file_name="decoder_model_merged.onnx",
        )
        self.forced_decoder_ids = self.processor.get_decoder_prompt_ids(
            language="indonesian", task="transcribe"
        )
        logger.info("Whisper ONNX loaded ✓")

        # ── Silero VAD ──────────────────────────────────────────────────
        self.vad_model, self.vad_utils = torch.hub.load(
            repo_or_dir="snakers4/silero-vad",
            model="silero_vad",
            force_reload=False,
            trust_repo=True,
        )
        (self.get_speech_timestamps, _, self.read_audio, *_) = self.vad_utils
        logger.info("Silero VAD loaded ✓")

        self.rolling_buffer = np.array([], dtype=np.float32)
        self.last_text      = ""
        self.is_active      = True
        self.lock           = threading.Lock()

        logger.info("ASRPipeline (Whisper ONNX + Groq hybrid) siap ✓")

    # ── helpers ─────────────────────────────────────────────────────────

    def stop(self):
        self.rolling_buffer = np.array([], dtype=np.float32)
        self.last_text      = ""
        self.is_active      = True
        logger.info("Pipeline stopped ✓")

    def preprocess(self, audio: np.ndarray) -> np.ndarray:
        if len(audio) > SAMPLE_RATE * 0.1:
            try:
                audio = nr.reduce_noise(y=audio, sr=SAMPLE_RATE, stationary=True)
            except Exception:
                pass
        return audio.astype(np.float32)

    def get_vad_confidence(self, audio: np.ndarray) -> float:
        VAD_CHUNK   = 512
        confidences = []
        for i in range(0, len(audio) - VAD_CHUNK + 1, VAD_CHUNK):
            chunk  = audio[i : i + VAD_CHUNK]
            tensor = torch.from_numpy(chunk)
            conf   = self.vad_model(tensor, SAMPLE_RATE).item()
            confidences.append(conf)
        return max(confidences) if confidences else 0.0

    def has_enough_energy(self, audio: np.ndarray) -> bool:
        return np.sqrt(np.mean(audio ** 2)) > 0.01

    def is_valid_text(self, text: str) -> bool:
        t = text.lower().strip()
        if len(t) < 3:
            return False
        if t in EXACT_BLACKLIST:
            return False
        for bl in BLACKLIST:
            if t.startswith(bl):
                return False
        for s in ENGLISH_STARTERS:
            if t.startswith(s):
                return False
        words = t.replace(".", "").replace(",", "").replace("!", "").split()
        if len(words) >= 3:
            mc = max(set(words), key=words.count)
            if words.count(mc) / len(words) > 0.4:
                return False
        return True

    # ── primary: Whisper ONNX ───────────────────────────────────────────

    def whisper_transcribe(self, audio: np.ndarray) -> tuple[str, float]:
        """
        Returns (text, avg_logprob).
        avg_logprob ~ 0 = confident, << 0 = uncertain.
        """
        try:
            audio = self.preprocess(audio)
            inputs = self.processor(
                audio,
                sampling_rate=SAMPLE_RATE,
                return_tensors="pt",
            )
            generated = self.onnx_model.generate(
                **inputs,
                forced_decoder_ids=self.forced_decoder_ids,
                return_dict_in_generate=True,
                output_scores=True,
            )
            text = self.processor.batch_decode(
                generated.sequences, skip_special_tokens=True
            )[0].strip()

            # Hitung avg_logprob dari scores
            if generated.scores:
                import torch
                log_probs = []
                for score, token_id in zip(
                    generated.scores,
                    generated.sequences[0][len(self.forced_decoder_ids) + 1:],
                ):
                    lp = torch.nn.functional.log_softmax(score[0], dim=-1)
                    log_probs.append(lp[token_id].item())
                avg_logprob = float(np.mean(log_probs)) if log_probs else -1.0
            else:
                avg_logprob = -1.0

            logger.info(f"[ONNX] '{text}' | logprob={avg_logprob:.3f}")
            return text, avg_logprob

        except Exception as e:
            logger.error(f"Whisper ONNX error: {e}")
            return "", -99.0

    # ── fallback: Groq ──────────────────────────────────────────────────

    def groq_transcribe(self, audio: np.ndarray) -> str:
        try:
            audio     = self.preprocess(audio)
            wav_bytes = audio_to_wav_bytes(audio)
            result    = [""]
            success   = [False]
            start     = time.time()

            def call_groq():
                try:
                    client = _get_groq_client()
                    transcription = client.audio.transcriptions.create(
                        file=("audio.wav", wav_bytes),
                        model="whisper-large-v3-turbo",
                        language="id",
                        response_format="text",
                        prompt="",
                    )
                    result[0]  = (
                        transcription.strip()
                        if isinstance(transcription, str)
                        else str(transcription).strip()
                    )
                    success[0] = True
                except Exception as e:
                    logger.error(f"Groq error: {e}")

            t = threading.Thread(target=call_groq)
            t.start()
            t.join(timeout=GROQ_TIMEOUT)

            elapsed = time.time() - start
            if success[0] and result[0]:
                logger.info(f"[Groq {elapsed:.2f}s] >> {result[0]}")
                return result[0]
            logger.warning(f"[Groq {elapsed:.2f}s] timeout/empty")
            return ""
        except Exception as e:
            logger.error(f"Groq exception: {e}")
            return ""

    # ── main entry ──────────────────────────────────────────────────────

    def transcribe_chunk(self, audio_chunk: np.ndarray) -> str | None:
        with self.lock:
            if not self.is_active:
                return None

            # VAD check
            vad_conf = self.get_vad_confidence(audio_chunk)
            logger.debug(f"VAD conf: {vad_conf:.3f}")
            if vad_conf <= VAD_THRESHOLD:
                self.rolling_buffer = np.array([], dtype=np.float32)
                return None

            if not self.has_enough_energy(audio_chunk):
                return None

            # Accumulate buffer
            self.rolling_buffer = np.concatenate([self.rolling_buffer, audio_chunk])
            if len(self.rolling_buffer) > BUFFER_SAMPLES:
                self.rolling_buffer = self.rolling_buffer[-BUFFER_SAMPLES:]

            if len(self.rolling_buffer) < BUFFER_SAMPLES:
                return None

            audio_to_process = self.rolling_buffer.copy()
            self.rolling_buffer = self.rolling_buffer[BUFFER_SAMPLES // 2:]

            # ── Hybrid logic ────────────────────────────────────────────
            text, avg_logprob = self.whisper_transcribe(audio_to_process)

            if not text or avg_logprob < CONFIDENCE_THRESHOLD:
                # Confidence rendah → fallback Groq
                source = "groq"
                text   = self.groq_transcribe(audio_to_process)
            else:
                source = "onnx"

            logger.info(f"[{source}] final: '{text}'")

            if not text:
                return None
            if not self.is_valid_text(text):
                logger.info(f"Filtered: {text}")
                return None

            self.last_text = text
            return text.strip()