from typing import TypedDict, Any, Optional, Tuple
from langgraph.graph import StateGraph, END

from langfuse import get_client
from agent.stage.unified_agent import UnifiedAgent
from agent.stage.unified_state import ConversationPhase
from agent.utils.utils import (
    process_mailing,
    process_plan,
    get_gender,
    get_plan,
    normalize_name,
    sentence_confidence,
    start_message_argumentation,
)

from agent.classifier.check_guardrail import CheckGuardrail
from agent.classifier.conversation_summary import ConversationSummary
from agent.classifier.conversation_classifier import (
    get_historic_conversation,
    ConversationClassifier,
    format_success_purchase,
    format_call,
)

from app.common.timed import timed
from langfuse.langchain import CallbackHandler
import asyncio
import copy
import os
import json
import math
from app.utils.logging import setup_minimal_logging

logger = setup_minimal_logging()
langfuse = get_client()
langfuse_handler = CallbackHandler()


class PipelineState(TypedDict, total=False):
    stage: str
    user_input: dict
    auto: bool
    output: str
    is_attack: bool


class CustomerPipelineUnified:
    """Pipeline unificado: um único agente com prompt dinâmico para toda a conversa."""

    def __init__(self, mailing, streaming=False):
        self.streaming = streaming
        self.stage = "presentation"
        self.classified_conversation = None
        self.accepted_purchase = False
        self.name = None
        self.attacks = 0

        self.prompt_vars = {
            "raw_customer": mailing,
            "customer": None,
            "actual_plan": None,
            "target_plan": None,
            "elegibility": None,
        }

        self.trace_id = langfuse.create_trace_id()
        self.intro = None
        self.graph = None
        self.agent: UnifiedAgent | None = None

        # Buffer do último output (aguarda sinal do backend para subir no Langfuse)
        self.pending_input = None
        self.pending_output = None
        self.pending_stage = None
        self.pending_metadata = None
        print("RODANDO O PIPELINE UNIFICADO")
    # ====================================================================================
    # Buffer + commit no Langfuse
    # ====================================================================================

    def _buffer_langfuse_trace(self, input_value: Any, output_value: str):
        """
        Guarda o último par (input, output) gerado pelo agente para só subir no Langfuse
        quando o backend sinalizar se foi ouvido por completo ou interrompido.
        """
        self.pending_input = input_value
        self.pending_output = output_value
        self.pending_stage = self.stage
        self.pending_metadata = self.get_formatted_metadata()

    def update_langfuse_trace(self, final_output: str, was_interrupted: bool = False):
        """
        Faz o commit no Langfuse usando os campos pendentes (input/stage/metadata)
        e o output final (completo ou truncado).
        """
        if self.pending_stage is None:
            return

        metadata = self.pending_metadata or {}
        metadata = {**metadata, "interrupted": was_interrupted}

        with langfuse.start_as_current_observation(
            as_type="span",
            name=self.pending_stage,
            trace_context={"trace_id": self.trace_id},
        ) as span:
            langfuse.update_current_span(
                metadata=metadata,
                name=self.pending_stage,
                output=final_output,
                input=self.pending_input,
            )

    def update_langfuse_auto(self, input, output):
        with langfuse.start_as_current_observation(
                    as_type="span",
                    name=self.stage,
                    trace_context={"trace_id": self.trace_id},
                ) as span:
            langfuse.update_current_span(
                        metadata=self.get_formatted_metadata(),
                        name=self.stage,
                        output=output,
                        input=input
                    )

    # ====================================================================================
    # Backend sinaliza consumo/interrupção (TTS)
    # ====================================================================================

    def set_interruption(
        self,
        was_interrupted: bool,
        listened_text: Optional[str] = None,
        skipped_vacalization: bool = False):
        
        logger.info(
            "SET_INTERRUPTION | was_interrupted=%s | listened_text=%r | skipped_vacalization=%s",
            was_interrupted,
            listened_text,
            skipped_vacalization,
        )

        if not self.pending_output:
            return

        # Definimos a marca de interrupção
        INTERRUPTION_TAG = "###interrupção###"

        if skipped_vacalization:
            # Caso 1: O áudio nem chegou a tocar (interrupção imediata ou erro)
            final_output = INTERRUPTION_TAG
            was_interrupted_flag = True 
            self.agent.state.restore_state()
            self.agent.set_interruption_message(final_output)

        elif was_interrupted:
            # Caso 2: Estava falando e foi cortado
            # Se ouviu algo, usa o trecho + tag. Se não ouviu nada, apenas a tag.
            final_output = f"{listened_text or ''} {INTERRUPTION_TAG}".strip()
            was_interrupted_flag = True
            self.agent.set_interruption_message(final_output)
            
        else:
            # Caso 3: Fluxo normal (não foi interrompido)
            final_output = self.pending_output
            was_interrupted_flag = False

        # Atualiza o Rastreamento (Langfuse)
        self.update_langfuse_trace(
            final_output,
            was_interrupted=was_interrupted_flag,
        )

        # Limpa buffer
        self.pending_input = None
        self.pending_output = None
        self.pending_stage = None
        self.pending_metadata = None

    # ====================================================================================
    # Start + Prepare
    # ====================================================================================

    def start(self):
        first_name = self.prompt_vars["raw_customer"]["NOME_CLIENTE_COMPLETO"].split()[0].lower()
        self.intro = (
            f"Olá, meu nome é Helena, sou consultora de vendas da TIM, para sua segurança a ligação está sendo gravada. "
            f"Eu estou falando com {first_name}?"
        )
        self.stage = "presentation"
        self.name = "presentation"
        return self.stage, self.intro

    def _generate_argumentation_start_text(self):
        """Gera a mensagem fixa de abertura da argumentação."""
        text = start_message_argumentation({
            "cliente_alvo_primeiro_nome": self.prompt_vars["customer"]["cliente_alvo_primeiro_nome"].lower(),
            "meses_restantes_fidelizacao": self.prompt_vars["actual_plan"]["meses_restantes_fidelizacao"],
            "quanto_pago_a_mais": self.prompt_vars["target_plan"]["quanto_pago_a_mais"],
            "quanto_pago_a_mais_int": abs(math.ceil(
                self.prompt_vars["raw_customer"]["VLR_FINAL_PLANO_DEST"] - self.prompt_vars["raw_customer"]["MEDIA_RECARGA_2"]
            )),
            "gb_plano_atual": self.prompt_vars["actual_plan"]["gb_plano_atual"],
            "valor_plano_sem_fidelização": self.prompt_vars["actual_plan"]["valor_plano_sem_fidelização"],
            "dados_GB": self.prompt_vars["target_plan"]["dados_GB"],
            "beneficios_001": self.prompt_vars["target_plan"]["beneficios_001"],
            "valor_plano_final": self.prompt_vars["target_plan"]["valor_plano_final"],
            "gb_alvo_diferenca": self.prompt_vars["target_plan"]["gb_alvo_diferenca"],
            "preco_por_dia": self.prompt_vars["target_plan"]["preço_por_dia"]
        })
        return text

    def _generate_formalization_start_text(self):
        """Gera o texto obrigatório e literal de abertura da formalização."""
        customer = self.prompt_vars["customer"]
        target = self.prompt_vars["target_plan"]
        tratamento = customer.get("TRATAMENTO", "Você")

        text = (
            f"Pra finalizar eu preciso apresentar uma série de informações sobre o novo plano. "
            f"Vou começar: {tratamento} está adquirindo uma oferta do plano {target['plano']}, "
            f"que por doze meses terá o custo de {target['valor_plano_final']}. "
            f"Este é um desconto de {target['valor_desconto']} em relação ao valor total do plano, "
            f"que é {target['valor_plano_bruto']}. Nesta oferta, {tratamento} tem {target['dados_GB']} "
            f"gigas para navegar na internet. Além disso, você tem os seguintes benefícios no novo plano: "
            f"{target['beneficios_todos']} e também recebe uma série de serviços de valor agregado. "
            f"Você pode conhecê-los no APP meu Tim. A sua data de vencimento e forma de pagamento "
            f"permanecem os mesmos. Na sua próxima fatura será cobrado o valor integral do plano antigo, "
            f"mais o proporcional dos dias utilizados do plano adquirido. {tratamento} confirma a migração "
            f"para o {target['plano']} no valor de {target['valor_plano_final']} fidelizado por doze meses?? "
            f"Se sim, diga eu confirmo."
        )
        return text

    def _generate_data_confirmation_start_text(self):
        """Gera o texto de abertura da confirmação de dados."""
        nome = self.prompt_vars["customer"]["NOME_CLIENTE_COMPLETO"]
        text = (
            f"Ótimo! então vamos ativar agora para você já aproveitar! "
            f"Para prosseguirmos você pode por favor confirmar se seu nome é {nome}? "
            f"Se sim, diga eu confirmo"
        )
        return text

    # ====================================================================================
    # Prepare — instancia o UnifiedAgent
    # ====================================================================================

    def prepare(self, elegibility: bool, protocol: str):
        self.stage = "presentation"
        self.name = "presentation"

        self.prompt_vars["customer"] = copy.deepcopy(self.prompt_vars["raw_customer"])
        self.prompt_vars["customer"]["PROTOCOLO"] = protocol

        self.prompt_vars["elegibility"] = elegibility

        self.prompt_vars["actual_plan"] = process_plan({}, self.prompt_vars["customer"], "actual")
        self.prompt_vars["target_plan"] = process_plan(
            get_plan(self.prompt_vars["customer"]["PLANO_DESTINO"]),
            self.prompt_vars["customer"],
            "target",
        )

        self.prompt_vars["customer"] = get_gender(self.prompt_vars["customer"])
        self.prompt_vars["customer"]["NOME_CLIENTE_COMPLETO"] = normalize_name(
            self.prompt_vars["customer"]["NOME_CLIENTE_COMPLETO"]
        )

        self.prompt_vars["customer"] = process_mailing(self.prompt_vars["customer"])

        # ── Instancia o agente unificado ──
        self.agent = UnifiedAgent(
            expected_cpf_last_3=self.prompt_vars["customer"]["NUM_CPF_CNPJ_CLIENTE"],
            expected_birth_date=self.prompt_vars["customer"].get("DAT_NASCIMENTO", ""),
            streaming=self.streaming,
            prompt_vars={
                "customer_first_name": self.prompt_vars["customer"]["NOME_CLIENTE_COMPLETO"]
                    .split()[0].lower(),
                "customer_full_name": self.prompt_vars["customer"]["NOME_CLIENTE_COMPLETO"],
                "customer": self.prompt_vars["customer"],
                "actual_plan": self.prompt_vars["actual_plan"],
                "target_plan": self.prompt_vars["target_plan"],
                "tratamento": self.prompt_vars["customer"].get("TRATAMENTO", "O senhor"),
                "plano_target": self.prompt_vars["target_plan"]["plano"],
                "valor_plano_final": self.prompt_vars["target_plan"]["valor_plano_final"],
                "valor_desconto": self.prompt_vars["target_plan"]["valor_desconto"],
                "valor_plano_bruto": self.prompt_vars["target_plan"]["valor_plano_bruto"],
                "dados_GB": self.prompt_vars["target_plan"]["dados_GB"],
                "beneficios_todos": self.prompt_vars["target_plan"]["beneficios_todos"],
                "elegibility": self.prompt_vars["elegibility"],
            },
        )

        # Injeta o início da conversa no histórico
        self.agent.inject_user_message("###START###")
        self.agent.inject_ai_message(self.intro)

        self.check_guardrail = CheckGuardrail()
        self.conversation_classifier = ConversationClassifier()
        self.conversation_summary = ConversationSummary()

        self.update_langfuse_auto(
            {"text": "###START###", "prompt_vars": self.prompt_vars},
            self.intro,
        )

        self.graph = self._build_graph()

    # ====================================================================================
    # Grafo simplificado: guardrail → unified → END
    # ====================================================================================

    def _build_graph(self):
        g = StateGraph(PipelineState)

        g.add_node("guardrail", self._lg_guardrail)
        g.add_node("unified", self._lg_unified)

        g.set_entry_point("guardrail")

        g.add_conditional_edges(
            "guardrail",
            self._lg_route_from_guardrail,
            {
                "unified": "unified",
                "end": END,
            },
        )

        g.add_conditional_edges(
            "unified",
            self._lg_route_after_stage,
            {
                "unified": "unified",
                "end": END,
            },
        )

        return g.compile()

    # ====================================================================================
    # STT normalization
    # ====================================================================================

    def _normalize_stt_input(self, user_input: Any) -> Tuple[str, Optional[list]]:
        """
        api_text  -> retorna (texto, None)
        raw_json  -> retorna (data.text, data.words)
        """
        mode = os.getenv("STT_OUTPUT_MODE", "api_text").lower()

        if mode != "raw_json":
            return str(user_input or ""), None

        try:
            payload = json.loads(user_input) if isinstance(user_input, str) else user_input
            data = (payload or {}).get("data", {})
            text = str(data.get("text", "") or "")
            words = data.get("words", None)
            return text, words
        except Exception:
            return str(user_input or ""), None

    # ====================================================================================
    # Guardrail (idêntico ao original)
    # ====================================================================================

    def _lg_guardrail(self, state: PipelineState) -> PipelineState:
        raw_input = state.get("user_input", "")
        user_input, words = self._normalize_stt_input(raw_input)

        # Guardrail de confiança (somente raw_json, sem limite)
        if words is not None:
            conf = sentence_confidence(words)
            print("confiança", conf)
            if conf < 0.0 or user_input in [None, ""]:
                output = "Desculpe, não consegui entender bem. Pode repetir mais devagar, por favor?"

                self._buffer_langfuse_trace(user_input, output)

                return {
                    "is_attack": True,
                    "output": output,
                    "auto": False,
                    "stage": state.get("stage", self.stage),
                    "user_input": user_input,
                }

        # Guardrail original (ataque)
        is_attack = self._check_guardrail(user_input)
        output = "Desculpe, não entendi. Poderia repetir?"

        if self.attacks > 4:
            output = (
                "Notei que temos um problema de comunicação. "
                "Irei finalizar a chamada. A TIM agradece a sua atenção"
            )

            self._buffer_langfuse_trace(user_input, output)

            return {
                "is_attack": False,
                "output": output,
                "auto": False,
                "stage": "DONE",
                "user_input": user_input,
            }

        if is_attack:
            self.attacks += 1

            self._buffer_langfuse_trace(user_input, output)

            return {
                "is_attack": True,
                "output": output,
                "auto": False,
                "stage": state.get("stage", self.stage),
                "user_input": user_input,
            }

        return {"is_attack": False, "user_input": user_input}

    def _lg_route_from_guardrail(self, state: PipelineState) -> str:
        if state.get("is_attack"):
            return "end"

        stage = state.get("stage") or self.stage
        if stage == "DONE":
            return "end"

        return "unified"

    def _lg_route_after_stage(self, state: PipelineState) -> str:
        if state.get("stage") == "DONE":
            return "end"

        if state.get("auto"):
            return "unified"

        return "end"

    # ====================================================================================
    # Nó unificado — substitui os 4 nós separados
    # ====================================================================================

    def _lg_unified(self, state: PipelineState) -> PipelineState:
        """Nó único que delega ao UnifiedAgent e gerencia transições automáticas."""
        self.stage = state.get("stage", self.stage)
        self.name = self.agent.state.current_phase.value

        # ── Se é auto-start de uma nova fase, injeta mensagem fixa ──
        if state.get("auto"):
            phase = self.agent.state.current_phase

            if phase == ConversationPhase.ARGUMENTATION:
                # Mensagem fixa de argumentação
                text_arg = self._generate_argumentation_start_text()
                self.agent.inject_user_message("###START###")
                self.agent.inject_ai_message(text_arg)

                # Buffer para Langfuse (transição automática)
                self._buffer_langfuse_trace("###START###", text_arg)

                # Retorna a mensagem fixa como output — cliente ouve isso
                return {
                    "stage": phase.value,
                    "output": text_arg,
                    "auto": False,
                }

            elif phase == ConversationPhase.DATA_CONFIRMATION:
                # Roda com ###START### para gerar a mensagem de confirmação
                text_dc = self._generate_data_confirmation_start_text()
                self.agent.inject_user_message("###START###")
                self.agent.inject_ai_message(text_dc)

                self._buffer_langfuse_trace("###START###", text_dc)

                return {
                    "stage": phase.value,
                    "output": text_dc,
                    "auto": False,
                }

            elif phase == ConversationPhase.FORMALIZATION:
                # Texto obrigatório e literal da formalização
                text_form = self._generate_formalization_start_text()
                self.agent.inject_user_message("###START###")
                self.agent.inject_ai_message(text_form)

                self._buffer_langfuse_trace("###START###", text_form)

                return {
                    "stage": phase.value,
                    "output": text_form,
                    "auto": False,
                }

        # ── Fluxo normal: roda o agente unificado ──
        text = state.get("user_input", "")
        result = self.agent.run_unified(text)

        output = result["output"]
        new_phase = result["phase"]
        auto = result["auto"]

        # Buffer para Langfuse
        self._buffer_langfuse_trace(text, output)

        # Atualiza stage do pipeline
        if new_phase == "done":
            self.stage = "DONE"
            # Verifica se foi aceite final da formalização
            if self.agent.state.form_accepted:
                self.accepted_purchase = True
            return {"stage": "DONE", "output": output, "auto": False}

        self.stage = new_phase

        # Se houve transição automática, loga no Langfuse antes de trocar
        if auto:
            self.update_langfuse_auto(text, output)

        return {
            "stage": new_phase,
            "output": output,
            "auto": auto,
        }

    # ====================================================================================
    # Public API
    # ====================================================================================

    @timed("Agent run")
    def run(self, user_input: Any):
        if self.graph is None:
            raise RuntimeError("prepare() deve ser chamado antes de run().")

        if self.stage == "DONE":
            return "DONE", ""

        initial_state: PipelineState = {
            "stage": self.stage,
            "user_input": user_input,
            "auto": False,
        }

        out: PipelineState = self.graph.invoke(initial_state)

        self.stage = out.get("stage", self.stage)
        return self.stage, out.get("output", "")

    def _check_guardrail(self, user_input: str):
        return self.check_guardrail.run(user_input)["is_attack"]

    # ====================================================================================
    # end_service (mantém compatibilidade com _end_service)
    # ====================================================================================

    @timed("Agent end service")
    async def end_service(self):
        return await self._end_service()

    async def _end_service(self):
        if self.pending_output:
            self.set_interruption(was_interrupted=False)

        await asyncio.sleep(10)
        historical_conversation = get_historic_conversation(self.trace_id)

        self.classified_conversation = self.conversation_classifier.run(historical_conversation)
        self.summary = self.conversation_summary.run(historical_conversation)

        success_purchase = {}

        if self.accepted_purchase:
            success_purchase = format_success_purchase(self.prompt_vars["raw_customer"])

        formatted_call = format_call(self.prompt_vars["raw_customer"], self.classified_conversation, self.summary)
        output = {
            "success_purchase": success_purchase,
            "formatted_call": formatted_call,
        }

        with langfuse.start_as_current_observation(
            as_type="span",
            name=self.name,
            trace_context={"trace_id": self.trace_id},
        ) as span:
            langfuse.update_current_span(
                metadata={
                    "stage": "classify_conversation",
                    "classification": formatted_call["MOTIVO_STATUS"],
                },
                output=output,
                name=self.name,
                input=historical_conversation,
            )

            span.score_trace(
                name="success_purchase",
                value=1 if self.accepted_purchase else 0,
                data_type="BOOLEAN",
            )

        return output

    # ====================================================================================
    # Metadata — lê do UnifiedState centralizado
    # ====================================================================================

    def get_formatted_metadata(self):
        if self.agent is None:
            return {"stage": self.stage, "attacks": self.attacks}

        state = self.agent.state
        phase = state.current_phase.value

        base = {
            "stage": phase,
            "attacks": self.attacks,
            "current_phase": phase,
            "end_conversation": state.end_conversation,
        }

        if phase == "presentation":
            base.update({
                "presentation.name_validation_attempt": state.name_validation_attempt,
                "presentation.eligibility": self.prompt_vars.get("elegibility"),
                "presentation.is_target_customer": state.is_target_customer,
            })

        elif phase == "argumentation":
            base.update({
                "argumentation.accepted": state.arg_accepted,
                "argumentation.accepted_count": state.arg_accepted_count,
                "argumentation.declines": state.arg_declines,
                "argumentation.off_topic": state.arg_off_topic,
            })

        elif phase == "data_confirmation":
            base.update({
                "data_confirmation.authenticated": state.is_authenticated,
                "data_confirmation.validated_target_customer": state.validated_target_customer,
                "data_confirmation.validated_cpf": state.validated_cpf,
                "data_confirmation.validated_birth_date": state.validated_birth_date,
                "data_confirmation.attempt_customer": state.attempt_customer,
                "data_confirmation.attempt_cpf": state.attempt_cpf,
                "data_confirmation.attempt_birth_date": state.attempt_birth_date,
            })

        elif phase == "formalization":
            base.update({
                "formalization.accepted": state.form_accepted,
                "formalization.declines": state.form_declines,
                "formalization.questions": state.form_questions,
            })

        return base
