# app/utils/interruption_tracker.py
from __future__ import annotations

import inspect
import logging
import re
import unicodedata
from dataclasses import dataclass
from typing import Any, Callable, Dict, Optional, Set
"""from agent.classifier.interruption_classifier import InterruptionClassifier

interruption_classifier = InterruptionClassifier()

def is_backchannel(text: str) -> bool:
    result = interruption_classifier.run(text)
    print(f"[is_backchannel] {text} -> {result}")
    return result"""

_BACKCHANNEL_SET = {
    "ta", "tá", "ok", "okay", "certo", "beleza", "blz",
    "tudo bem", "tudo bem?", "tudo bom", "tudo bom?",
    "aham", "uhum", "hum", "hm", "uh-huh",
    "entendi", "isso", "claro",
}


def _norm_pt(s: str) -> str:
    s = (s or "").strip().lower()
    s = unicodedata.normalize("NFKD", s)
    s = "".join(ch for ch in s if not unicodedata.combining(ch))
    s = re.sub(r"\s+", " ", s)
    s = re.sub(r"[^\w\s?]", "", s)  # remove pontuação (mantém ? e -)
    return s.strip()


def is_backchannel(text: str) -> bool:
    t = _norm_pt(text)
    print(f"[is_backchannel] {text} -> {t}")
    if not t:
        return True
    if t in _BACKCHANNEL_SET:
        return True
    if len(t) <= 2:
        return True
    #if t.startswith("ta ") or t.startswith("tá "):
    #    return True
    return False

def remaining_after_prefix(full: str, spoken: str) -> str:
    """
    Mais robusto:
      - tenta prefixo exato
      - tenta common-prefix (char a char)
      - tenta find()
      - fallback: retorna full (melhor do que ficar em silêncio)
    """
    full = full or ""
    spoken = spoken or ""
    if not full:
        return ""

    if not spoken:
        return full

    if full.startswith(spoken):
        return full[len(spoken):]

    # longest common prefix
    m = 0
    for a, b in zip(full, spoken):
        if a == b:
            m += 1
        else:
            break
    if m > 0:
        return full[m:]

    idx = full.find(spoken)
    if idx >= 0:
        return full[idx + len(spoken):]

    # fallback: repete o full (melhor do que silêncio)
    return full


@dataclass
class InterruptionTrace:
    stage: str
    full_text: str
    started_at: float
    spoken_text: str = ""
    ended_at: float = 0.0
    interrupted: bool = False
    interrupter_text: str = ""
    interrupter_transcript_raw: str = ""


class InterruptionTracker:
    """
    Encapsula:
      - detecção de backchannel
      - rastreio de interrupção (full vs spoken)
      - captura do que o usuário disse ao interromper
      - auto-resume se for backchannel OU interrupção vazia
      - payload "interrupt" para ser passado ao pipeline.run
      - interrupt_now() para cortar o áudio IMEDIATAMENTE (sem esperar is_final)
    """

    def __init__(
        self,
        *,
        allow_interruptions_for_stage: Callable[[str], bool],
        logger: Optional[logging.Logger] = None,
        auto_resume_stages: Optional[Set[str]] = None,
        max_auto_resume: int = 2,
        auto_resume_on_empty_interrupt: bool = True,
        auto_resume_backchannel_any_stage: bool = True,
        auto_resume_empty_any_stage: bool = True,
        on_interrupt: Optional[Callable[[dict], None]] = None,
    ) -> None:
        self.allow_interruptions_for_stage = allow_interruptions_for_stage
        self.log = logger or logging.getLogger(__name__)
        self.on_interrupt = on_interrupt

        self.auto_resume_stages = {s.upper() for s in (auto_resume_stages or {"INTRO", "PRESENTATION"})}
        self.max_auto_resume = max(0, int(max_auto_resume))
        self.auto_resume_on_empty_interrupt = bool(auto_resume_on_empty_interrupt)

        self.auto_resume_backchannel_any_stage = bool(auto_resume_backchannel_any_stage)
        self.auto_resume_empty_any_stage = bool(auto_resume_empty_any_stage)

        self._loop = None
        self._agent_speaking = None
        self._session = None

        self.active_trace: Optional[InterruptionTrace] = None
        self.last_interrupt_info: Optional[dict] = None

        # ✅ handle do say() atual (para cortar imediatamente)
        self._current_handle = None

        # ✅ buffer caso note_user_interrupt chegue antes do active_trace
        self._pending_interrupter_text: str = ""
        self._pending_interrupter_raw: str = ""

    def is_backchannel_text(self, text: str) -> bool:
        return is_backchannel(text)

    def attach(self, session: Any, agent_speaking_event: Any, loop: Any) -> None:
        self._loop = loop
        self._agent_speaking = agent_speaking_event
        self._session = session
        session.on("conversation_item_added")(self._on_conversation_item_added)

    def consume_last_interrupt_info(self) -> Optional[dict]:
        info = self.last_interrupt_info
        self.last_interrupt_info = None
        return info

    def note_user_interrupt(self, user_text: str, raw_transcript: str) -> None:
        """
        Registra o texto do usuário que interrompeu.

        ✅ Se ainda não existe active_trace (race com callbacks),
        guarda num buffer e aplica assim que começar o próximo say_with_trace.
        """
        if not user_text:
            return

        if self.active_trace and not self.active_trace.interrupter_text:
            self.active_trace.interrupter_text = user_text
            self.active_trace.interrupter_transcript_raw = raw_transcript or ""
            return

        # fallback: guarda para aplicar quando o trace começar
        if not self._pending_interrupter_text:
            self._pending_interrupter_text = user_text
            self._pending_interrupter_raw = raw_transcript or ""

    def should_ignore_user_input(self, user_text: str, agent_is_speaking: bool) -> bool:
        if not user_text:
            return True
        if agent_is_speaking and is_backchannel(user_text):
            return True
        return False

    def interrupt_now(self) -> bool:
        """
        Tenta interromper a fala atual imediatamente.
        Retorna True se conseguiu disparar algum método de interrupção.
        """
        h = self._current_handle

        # 1) tenta no handle retornado por say()
        if h is not None:
            for meth in ("interrupt", "cancel", "stop"):
                fn = getattr(h, meth, None)
                if callable(fn):
                    try:
                        ret = fn()
                        if inspect.isawaitable(ret) and self._loop:
                            self._loop.create_task(ret)
                        return True
                    except Exception:
                        self.log.exception("interrupt_now falhou usando handle.%s()", meth)
                        return False

        # 2) fallback: tenta no session (se existir)
        s = self._session
        if s is not None:
            for meth in ("interrupt", "cancel", "stop"):
                fn = getattr(s, meth, None)
                if callable(fn):
                    try:
                        ret = fn()
                        if inspect.isawaitable(ret) and self._loop:
                            self._loop.create_task(ret)
                        return True
                    except TypeError:
                        continue
                    except Exception:
                        self.log.exception("interrupt_now falhou usando session.%s()", meth)
                        return False

        return False

    def _on_conversation_item_added(self, ev: Any) -> None:
        """
        Captura o item do assistant "commitado".
        Quando há interrupção, o texto vem truncado e item.interrupted=True.
        """
        if not self.active_trace:
            return

        item = getattr(ev, "item", None)
        if not item:
            return

        role = getattr(item, "role", None)
        role_s = str(role).lower()
        if "assistant" not in role_s:
            return

        try:
            self.active_trace.spoken_text = getattr(item, "text_content", "") or ""
            self.active_trace.interrupted = bool(getattr(item, "interrupted", False))
        except Exception:
            pass

    def force_interrupt_on_end(self, *, reason: str = "call_end") -> Optional[dict]:
        """
        Se a call cair enquanto o agente está falando, pode não chegar um item final/truncado.
        Cria um payload de interrupção a partir do trace ativo e guarda em last_interrupt_info.
        """
        trace = self.active_trace
        if not trace:
            return None

        # Marca como interrompido por encerramento
        trace.interrupted = True

        spoken = trace.spoken_text or ""
        full = trace.full_text or ""
        remaining = remaining_after_prefix(full, spoken).strip()

        loop = self._loop
        now_ms = int(loop.time() * 1000) if loop else 0

        info = {
            "id": f"it-{now_ms}",
            "reason": reason,
            "stage": trace.stage,
            "assistant_full": full,
            "assistant_spoken": spoken,
            "assistant_remaining": remaining,
            "cut_chars": len(spoken),
            "ratio_spoken": (len(spoken) / max(1, len(full))) if full else 0.0,
            "elapsed_s": 0.0,
            "user_interrupt_text": trace.interrupter_text or "",
            "user_interrupt_transcript_raw": trace.interrupter_transcript_raw or "",
        }

        self.last_interrupt_info = info

        if self.on_interrupt:
            try:
                self.on_interrupt(info)
            except Exception:
                self.log.exception("on_interrupt falhou")

        return info

    async def _call_say(self, session: Any, text: str, allow: bool):
        res = session.say(text, allow_interruptions=allow, add_to_chat_ctx=True)
        if inspect.isawaitable(res):
            return await res
        return res

    async def say_with_trace(self, session: Any, stage: str, text_to_say: str) -> InterruptionTrace:
        stage_u = (stage or "UNKNOWN").upper()
        allow = bool(self.allow_interruptions_for_stage(stage_u))
        full = (text_to_say or "").strip()

        if not full:
            loop = self._loop
            t = loop.time() if loop else 0.0
            return InterruptionTrace(stage=stage_u, full_text="", started_at=t)

        attempt = 0
        last_trace: Optional[InterruptionTrace] = None

        while True:
            loop = self._loop
            start_t = loop.time() if loop else 0.0

            trace = InterruptionTrace(stage=stage_u, full_text=full, started_at=start_t)
            self.active_trace = trace

            # ✅ aplica interrupter pendente, caso tenha chegado antes do trace existir
            if self._pending_interrupter_text and not trace.interrupter_text:
                trace.interrupter_text = self._pending_interrupter_text
                trace.interrupter_transcript_raw = self._pending_interrupter_raw
                self._pending_interrupter_text = ""
                self._pending_interrupter_raw = ""

            if self._agent_speaking:
                self._agent_speaking.set()

            handle = None
            try:
                handle = await self._call_say(session, full, allow)

                # ✅ guarda handle atual para interrupt_now()
                self._current_handle = handle

                # ✅ aguarda playout
                join = getattr(handle, "join", None)
                if callable(join):
                    fut = join()
                    if fut is not None:
                        await fut
                else:
                    wait = getattr(handle, "wait_for_playout", None)
                    if callable(wait):
                        await wait()

            finally:
                end_t = loop.time() if loop else 0.0
                trace.ended_at = end_t

                if self._agent_speaking:
                    self._agent_speaking.clear()

                self.active_trace = None
                self._current_handle = None

            # interrupted flag
            trace.interrupted = bool(getattr(handle, "interrupted", False)) if handle else False

            # fallback spoken_text
            if trace.interrupted and not trace.spoken_text:
                trace.spoken_text = ""
            if (not trace.interrupted) and not trace.spoken_text:
                trace.spoken_text = trace.full_text

            last_trace = trace

            if trace.interrupted:
                remaining = remaining_after_prefix(trace.full_text, trace.spoken_text).strip()
                now_ms = int(loop.time() * 1000) if loop else 0

                info: Dict[str, Any] = {
                    "id": f"it-{now_ms}",
                    "reason": "barge_in",
                    "stage": trace.stage,
                    "assistant_full": trace.full_text,
                    "assistant_spoken": trace.spoken_text,
                    "assistant_remaining": remaining,
                    "cut_chars": len(trace.spoken_text or ""),
                    "ratio_spoken": (len(trace.spoken_text) / max(1, len(trace.full_text))) if trace.full_text else 0.0,
                    "elapsed_s": round(trace.ended_at - trace.started_at, 3),
                    "user_interrupt_text": trace.interrupter_text,
                    "user_interrupt_transcript_raw": trace.interrupter_transcript_raw,
                }

                self.last_interrupt_info = info

                if self.on_interrupt:
                    try:
                        self.on_interrupt(info)
                    except Exception:
                        self.log.exception("on_interrupt falhou")

                empty_interrupt = (not (trace.interrupter_text or "").strip()) and self.auto_resume_on_empty_interrupt
                back = is_backchannel(trace.interrupter_text)

                stage_ok = stage_u in self.auto_resume_stages
                if back and self.auto_resume_backchannel_any_stage:
                    stage_ok = True
                if empty_interrupt and self.auto_resume_empty_any_stage:
                    stage_ok = True

                can_resume = (
                    stage_ok
                    and attempt < self.max_auto_resume
                    and (empty_interrupt or back)
                )

                if can_resume and remaining:
                    attempt += 1
                    full = remaining
                    trace.interrupter_text = ""
                    trace.interrupter_transcript_raw = ""
                    continue

            break

        return last_trace or InterruptionTrace(stage=stage_u, full_text=full, started_at=(self._loop.time() if self._loop else 0.0))
