import json
from pathlib import Path
from typing import Optional, Dict, Any, Tuple, List

import numpy as np
from vosk import Model, KaldiRecognizer, SetLogLevel
SetLogLevel(-1)

DEFAULT_GRAMMAR = json.dumps([
    "sim",
    "não",
    "ok",
    "cancela",
    "cancelar",
    "ligar",
    "desligar",

    # +150 variações
    "claro",
    "claro que sim",
    "com certeza",
    "isso",
    "isso mesmo",
    "isso aí",
    "isso ai",
    "exato",
    "exatamente",
    "perfeito",
    "beleza",
    "certo",
    "certinho",
    "tá",
    "ta",
    "tá certo",
    "ta certo",
    "tá bom",
    "ta bom",
    "tudo bem",
    "pode",
    "pode sim",
    "pode ser",
    "pode ser sim",
    "pode deixar",
    "deixa",
    "de acordo",
    "concordo",
    "concordo sim",
    "confirmo",
    "eu confirmo",
    "confirmado",
    "positivo",
    "quero",
    "manda ver",
    "interessante",
    "legal",
    "bom",

    "negativo",
    "de jeito nenhum",
    "de forma alguma",
    "nem pensar",
    "não quero",
    "não mesmo",
    "não pode",
    "não dá",
    "não dá não",
    "não quero não",
    "quero não",
    "não quero isso",
    "não precisa",
    "não preciso",
    "dispenso",
    "para",
    "pare",
    "parar",
    "tá bom do jeito que tá",

    "tá ok",
    "ta ok",
    "tudo certo",
    "tudo ok",
    "certo então",
    "beleza então",
    "entendi",
    "entendido",
    "entendi sim",
    "entendido sim",
    "fechou",
    "fechado",
    "combinado",
    "show",
    "show de bola",
    "top",
    "maravilha",
    "tranquilo",
    "tranquila",
    "tá tranquilo",
    "ta tranquilo",

    "cancelamento",
    "cancele",
    "cancela aí",
    "cancela ai",
    "cancela agora",
    "cancelar agora",
    "pode cancelar",
    "pode cancelar aí",
    "pode cancelar ai",
    "quero cancelar",
    "eu quero cancelar",
    "quero o cancelamento",
    "faz o cancelamento",
    "fazer cancelamento",
    "solicitar cancelamento",
    "desisto",
    "desistir",
    "quero desistir",
    "para tudo",
    "pare tudo",
    "interrompe",
    "interromper",
    "encerra",
    "encerrar",
    "finaliza",
    "caro",

    "liga",
    "ligue",
    "liga aí",
    "liga ai",
    "liga agora",
    "ligar agora",
    "pode ligar",
    "pode ligar aí",
    "pode ligar ai",
    "pode ligar agora",
    "conecta",
    "conectar",
    "conecte",
    "inicia",
    "iniciar",
    "inicie",
    "ativar",
    "ativa",
    "ative",
    "habilitar",

    "desliga",
    "desligue",
    "desliga aí",
    "desliga ai",
    "desliga agora",
    "desligar agora",
    "pode desligar",
    "pode desligar aí",
    "pode desligar ai",
    "pode desligar agora",
    "desconecta",
    "desconectar",
    "desconecte",
    "encerra a ligação",
    "finaliza a ligação",
    "desativar",
    "desativa",
    "desative"
], ensure_ascii=False)


class VoskSTT:
    """
    STT com Vosk para CPU.

    Retorna SEMPRE:
      {"data": {"text": "...", "words": [{"word","start","end","probability"}]}}

    Observação:
      - Vosk retorna confiança por palavra em "conf" dentro de "result".
      - A gente converte "conf" -> "probability".
    """

    def __init__(
        self,
        model_path: str,
        sample_rate: int = 16000,
        grammar_json: Optional[str] = DEFAULT_GRAMMAR,
    ):
        mp = Path(model_path)
        if not mp.exists():
            raise FileNotFoundError(f"Vosk model_path não encontrado: {model_path}")

        self.sample_rate = sample_rate
        self.grammar_json = grammar_json
        self.model = Model(str(mp))

    def transcribe_pcm16le(self, pcm_bytes: bytes) -> Dict[str, Any]:
        """
        Recebe bytes PCM16LE mono (16k) e retorna dict no formato padrão.
        """
        if not pcm_bytes:
            return {"data": {"text": "", "words": []}}

        pcm16 = np.frombuffer(pcm_bytes, dtype=np.int16)  # zero-copy
        return self.recognize_pcm16(pcm16)

    def recognize_pcm16(self, pcm16: np.ndarray) -> Dict[str, Any]:
        pcm16 = np.asarray(pcm16)

        if pcm16.dtype != np.int16:
            raise TypeError(f"pcm16 must be int16, got {pcm16.dtype}")

        if pcm16.ndim == 2:
            if pcm16.shape[1] != 1:
                raise ValueError("Only mono supported; pass a single channel.")
            pcm16 = pcm16[:, 0]

        # usa grammar se configurado
        if self.grammar_json:
            rec = KaldiRecognizer(self.model, self.sample_rate, self.grammar_json)
        else:
            rec = KaldiRecognizer(self.model, self.sample_rate)

        rec.SetWords(True)
        rec.AcceptWaveform(pcm16.tobytes())
        raw = json.loads(rec.FinalResult())

        return self.to_data_words_format(raw)

    @staticmethod
    def to_data_words_format(raw: Dict[str, Any]) -> Dict[str, Any]:
        """
        Converte o retorno do Vosk (raw) para:
          {"data": {"text": "...", "words": [...]}}
        """
        if not isinstance(raw, dict):
            return {"data": {"text": str(raw), "words": []}}

        text = (raw.get("text") or "").strip()
        words_out: List[Dict[str, Any]] = []

        # formato padrão do Vosk com SetWords(True)
        # {"text":"...", "result":[{"word","start","end","conf"}]}
        if isinstance(raw.get("result"), list):
            for w in raw["result"]:
                if not isinstance(w, dict):
                    continue
                words_out.append({
                    "word": w.get("word", ""),
                    "start": float(w.get("start", 0.0) or 0.0),
                    "end": float(w.get("end", 0.0) or 0.0),
                    "probability": float(w.get("conf", w.get("probability", 0.0)) or 0.0),
                })
        # fallback: se vier "words" já no formato desejado
        elif isinstance(raw.get("words"), list):
            for w in raw["words"]:
                if not isinstance(w, dict):
                    continue
                words_out.append({
                    "word": w.get("word", ""),
                    "start": float(w.get("start", 0.0) or 0.0),
                    "end": float(w.get("end", 0.0) or 0.0),
                    "probability": float(w.get("probability", 0.0) or 0.0),
                })

        return {"data": {"text": text, "words": words_out}}

    @staticmethod
    def result_to_text(res: Dict[str, Any]) -> str:
        """
        Retorna apenas o texto do payload padronizado (ou do raw).
        """
        if not isinstance(res, dict):
            return str(res)
        if isinstance(res.get("data"), dict):
            t = res["data"].get("text")
            return t.strip() if isinstance(t, str) else ""
        t = res.get("text")
        return t.strip() if isinstance(t, str) else ""

    @staticmethod
    def result_to_text_and_payload(res: Dict[str, Any]) -> Tuple[str, Dict[str, Any]]:
        """
        Retorna (texto, payload_padronizado).
        """
        payload = VoskSTT.to_data_words_format(res) if not (isinstance(res, dict) and isinstance(res.get("data"), dict)) else res
        text = VoskSTT.result_to_text(payload)
        return text, payload
