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

from langfuse import get_client
from agent.stage.argumentation import ArgumentationAgent
from agent.stage.data_confirmation import DataConfirmationAgent
from agent.stage.formalization import FormalizationAgent
from agent.stage.presentation import PresentationAgent
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
import time
from datetime import datetime
from pathlib import Path
from app.utils.logging import setup_minimal_logging

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

# ===================== CONTROLE DE LOG DO AGENTE =====================
ENABLE_AGENT_LOG = True
# =====================================================================


class AgentLogger:
    """Logger que grava em TXT todo o ciclo de vida do pipeline do agente."""

    def __init__(self, telefone: str):
        self.telefone = str(telefone)
        self.start_ts = datetime.now()
        ts_str = self.start_ts.strftime("%H%M%S")

        base_dir = Path("log_agent") / self.start_ts.strftime("%Y-%m-%d") / self.start_ts.strftime("%H")
        base_dir.mkdir(parents=True, exist_ok=True)

        self.filepath = base_dir / f"{ts_str}_{self.telefone}.txt"
        self._lines: list[str] = []
        self._write_header()

    # ---- helpers internos ----
    def _write_header(self):
        self._append("=" * 80)
        self._append(f"AGENT LOG — Telefone: {self.telefone}")
        self._append(f"Início da sessão: {self.start_ts.isoformat()}")
        self._append("=" * 80)
        self._append("")

    def _ts(self) -> str:
        return datetime.now().strftime("%H:%M:%S.%f")[:-3]

    def _append(self, text: str):
        self._lines.append(text)

    def _flush(self):
        with open(self.filepath, "w", encoding="utf-8") as f:
            f.write("\n".join(self._lines))

    def _format_messages(self, messages: list) -> str:
        parts = []
        for msg in messages:
            role = type(msg).__name__.replace("Message", "").upper()
            content = getattr(msg, "content", str(msg))
            tool_calls = getattr(msg, "tool_calls", None)
            line = f"    [{role}] {content}"
            if tool_calls:
                line += f"  | tool_calls={tool_calls}"
            parts.append(line)
        return "\n".join(parts) if parts else "    (vazio)"

    def _format_state(self, state_obj) -> str:
        if state_obj is None:
            return "    (sem state)"
        d = {k: v for k, v in state_obj.__dict__.items() if not k.startswith("_")}
        parts = [f"    {k}: {v}" for k, v in d.items()]
        return "\n".join(parts)

    def _format_dict(self, d: dict, indent: int = 4) -> str:
        prefix = " " * indent
        parts = []
        for k, v in d.items():
            parts.append(f"{prefix}{k}: {v}")
        return "\n".join(parts) if parts else f"{prefix}(vazio)"

    # ---- métodos públicos de logging ----

    def log_section(self, title: str):
        self._append("")
        self._append("-" * 80)
        self._append(f"[{self._ts()}] {title}")
        self._append("-" * 80)
        self._flush()

    def log(self, label: str, content: str = ""):
        if content:
            self._append(f"  [{self._ts()}] {label}: {content}")
        else:
            self._append(f"  [{self._ts()}] {label}")
        self._flush()

    def log_system_prompt(self, agent_name: str, prompt: str):
        self._append(f"  [{self._ts()}] SYSTEM PROMPT ({agent_name}):")
        for line in prompt.splitlines():
            self._append(f"    | {line}")
        self._flush()

    def log_messages(self, label: str, messages: list):
        self._append(f"  [{self._ts()}] {label} — Mensagens ({len(messages)}):")
        self._append(self._format_messages(messages))
        self._flush()

    def log_agent_state(self, stage: str, state_obj):
        self._append(f"  [{self._ts()}] STATE ({stage}):")
        self._append(self._format_state(state_obj))
        self._flush()

    def log_metadata(self, metadata: dict):
        self._append(f"  [{self._ts()}] METADATA:")
        self._append(self._format_dict(metadata))
        self._flush()

    def log_stage_result(self, stage: str, output: str, next_stage: str, auto: bool):
        self._append(f"  [{self._ts()}] RESULTADO do estágio '{stage}':")
        self._append(f"    output: {output}")
        self._append(f"    próximo estágio: {next_stage}")
        self._append(f"    auto: {auto}")
        self._flush()


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


class CustomerPipeline:
    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 = None

        # Buffer do último output (aguarda sinal do backend)
        self.pending_input = None
        self.pending_output = None
        self.pending_stage = None
        self.pending_metadata = None

        # Logger do agente (TXT)
        telefone = mailing.get("NUM_TELEFONE", "sem_telefone")
        self.agent_log: AgentLogger | None = AgentLogger(telefone) if ENABLE_AGENT_LOG else None

    # ====================================================================================
    # 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).
        """
        #print('-'*30)
        #print(final_output,was_interrupted)
        #print(self.pending_input)
        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

    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

        # Log de interrupção
        if self.agent_log:
            self.agent_log.log_section("INTERRUPÇÃO")
            self.agent_log.log("was_interrupted", str(was_interrupted))
            self.agent_log.log("skipped_vacalization", str(skipped_vacalization))
            self.agent_log.log("listened_text", str(listened_text))
            self.agent_log.log("final_output", final_output)
            self.agent_log.log("was_interrupted_flag", str(was_interrupted_flag))
            self.agent_log.log_agent_state(self.stage, self.agent.state if self.agent else None)

        # 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


    # ====================================================================================
    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"

        if self.agent_log:
            self.agent_log.log_section("START")
            self.agent_log.log("stage", self.stage)
            self.agent_log.log("intro", self.intro)

        return self.stage, self.intro

    def argumentation_start(self):

        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"]
            })
        self.agent.inject_user_message("###START###")
        self.agent.inject_ai_message(text)
        return text

    # ====================================================================================
    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"])

        self.agent = PresentationAgent(
            streaming=self.streaming,
            prompt_vars={
                "customer_first_name": self.prompt_vars["customer"]["NOME_CLIENTE_COMPLETO"]
                .split()[0]
                .lower(),
                "elegibility": self.prompt_vars["elegibility"],
            },
        )

        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,)

        # Log do prepare
        if self.agent_log:
            self.agent_log.log_section("PREPARE")
            self.agent_log.log("elegibility", str(elegibility))
            self.agent_log.log("protocol", protocol)
            self.agent_log.log_system_prompt("presentation", self.agent.system_prompt)
            self.agent_log.log_messages("Mensagens iniciais (presentation)", self.agent.messages)
            self.agent_log.log("prompt_vars", str(self.prompt_vars))

        self.graph = self._build_graph()

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

        g.add_node("guardrail", self._lg_guardrail)
        g.add_node("presentation", self._lg_presentation)
        g.add_node("argumentation", self._lg_argumentation)
        g.add_node("data_confirmation", self._lg_data_confirmation)
        g.add_node("formalization", self._lg_formalization)

        g.set_entry_point("guardrail")

        g.add_conditional_edges(
            "guardrail",
            self._lg_route_from_guardrail,
            {
                "presentation": "presentation",
                "argumentation": "argumentation",
                "data_confirmation": "data_confirmation",
                "formalization": "formalization",
                "end": END,
            },
        )

        for node in ["presentation", "argumentation", "data_confirmation", "formalization"]:
            g.add_conditional_edges(
                node,
                self._lg_route_after_stage,
                {
                    "presentation": "presentation",
                    "argumentation": "argumentation",
                    "data_confirmation": "data_confirmation",
                    "formalization": "formalization",
                    "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

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

        if self.agent_log:
            self.agent_log.log_section(f"GUARDRAIL (estágio atual: {self.stage})")
            self.agent_log.log("user_input (normalizado)", user_input)
            self.agent_log.log("attacks acumulados", str(self.attacks))

        # 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?"

                if self.agent_log:
                    self.agent_log.log("guardrail_confiança", f"rejeitado (conf={conf})")
                    self.agent_log.log("output", output)

                # ANTES: subia direto
                # AGORA: buffer
                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)
        #print("is_attack", is_attack)
        is_attack = False
        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"
            )

            if self.agent_log:
                self.agent_log.log("guardrail", f"attacks > 4 — encerrando (attacks={self.attacks})")
                self.agent_log.log("output", output)

            # buffer
            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

            if self.agent_log:
                self.agent_log.log("guardrail", f"ataque detectado (attacks={self.attacks})")
                self.agent_log.log("output", output)

            # buffer
            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,
            }

        if self.agent_log:
            self.agent_log.log("guardrail", "passou (sem ataque)")

        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 stage

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

        if state.get("auto"):
            return state["stage"]

        return "end"

    # ====================================================================================
    # Stage nodes

    def _get_effective_input(self, state: PipelineState) -> str:
        return "###START###" if state.get("auto") else state.get("user_input", "")

    def _lg_presentation(self, state: PipelineState) -> PipelineState:
        self.stage = state.get("stage", self.stage)
        self.name = self.stage

        text = self._get_effective_input(state)

        if self.agent_log:
            self.agent_log.log_section("PRESENTATION")
            self.agent_log.log("input", text)
            self.agent_log.log_messages("Histórico ANTES do run", self.agent.messages)

        t0 = time.perf_counter()
        result = self.agent.run_presentation(text)
        elapsed = time.perf_counter() - t0

        if self.agent_log:
            self.agent_log.log("LLM tempo de execução", f"{elapsed:.3f}s")
            self.agent_log.log("LLM output", result["output"])
            self.agent_log.log("tool_calls", str(result.get("tools", [])))
            self.agent_log.log_agent_state("presentation", self.agent.state)
            self.agent_log.log_messages("Histórico DEPOIS do run", self.agent.messages)

        # buffer do output gerado
        self._buffer_langfuse_trace(text, result["output"])

        if self.agent.state.end_conversation:
            self.stage = "DONE"
            if self.agent_log:
                self.agent_log.log_stage_result("presentation", result["output"], "DONE", False)
            return {"stage": "DONE", "output": result["output"], "auto": False}

        if self.agent.state.is_target_customer:
            self.update_langfuse_auto(text, result["output"])
            self.stage = "argumentation"
            self.agent = ArgumentationAgent(
                streaming=self.streaming,
                prompt_vars=self.prompt_vars,
            )
            if self.agent_log:
                self.agent_log.log_stage_result("presentation", result["output"], "argumentation", True)
                self.agent_log.log_section("TRANSIÇÃO → ARGUMENTATION")
                self.agent_log.log_system_prompt("argumentation", self.agent.system_prompt)
            return {"stage": "argumentation", "auto": True}

        if not self.prompt_vars["elegibility"]:
            self.stage = "DONE"
            if self.agent_log:
                self.agent_log.log_stage_result("presentation", result["output"], "DONE (sem elegibilidade)", False)
            return {"stage": "DONE", "output": result["output"], "auto": False}

        if self.agent_log:
            self.agent_log.log_stage_result("presentation", result["output"], self.stage, False)

        return {"stage": self.stage, "output": result["output"], "auto": False}

    def _lg_argumentation(self, state: PipelineState) -> PipelineState:
        self.stage = state.get("stage", self.stage)
        self.name = self.stage

        if self.agent_log:
            self.agent_log.log_section("ARGUMENTATION")

        if state.get("auto"):
            result = {}
            result['output'] = self.argumentation_start()
            text = "###START###"
            if self.agent_log:
                self.agent_log.log("input (auto)", text)
                self.agent_log.log("argumentation_start output", result["output"])
        else:
            text = self._get_effective_input(state)

            if self.agent_log:
                self.agent_log.log("input", text)
                self.agent_log.log_messages("Histórico ANTES do run", self.agent.messages)

            t0 = time.perf_counter()
            result = self.agent.run_argumentation(text)
            elapsed = time.perf_counter() - t0

            if self.agent_log:
                self.agent_log.log("LLM tempo de execução", f"{elapsed:.3f}s")
                self.agent_log.log("LLM output", result["output"])
                self.agent_log.log("tool_calls", str(result.get("tools", [])))
                self.agent_log.log_messages("Histórico DEPOIS do run", self.agent.messages)

        if self.agent_log:
            self.agent_log.log_agent_state("argumentation", self.agent.state)

        # buffer do output gerado
        self._buffer_langfuse_trace(text, result["output"])

        if self.agent.state.accepted:
            self.update_langfuse_auto(text, result["output"])
            self.stage = "data_confirmation"

            self.agent = DataConfirmationAgent(
                expected_cpf_last_3=self.prompt_vars["customer"]["NUM_CPF_CNPJ_CLIENTE"],
                expected_birth_date=self.prompt_vars["customer"]["DAT_NASCIMENTO"],
                streaming=self.streaming,
                prompt_vars={"customer": self.prompt_vars["customer"]["NOME_CLIENTE_COMPLETO"]},
            )

            if self.agent_log:
                self.agent_log.log_stage_result("argumentation", result["output"], "data_confirmation", True)
                self.agent_log.log_section("TRANSIÇÃO → DATA_CONFIRMATION")
                self.agent_log.log_system_prompt("data_confirmation", self.agent.system_prompt)

            return {"stage": "data_confirmation", "auto": True}

        if self.agent.state.end_conversation:
            self.stage = "DONE"
            if self.agent_log:
                self.agent_log.log_stage_result("argumentation", result["output"], "DONE", False)
            return {"stage": "DONE", "output": result["output"], "auto": False}

        if self.agent_log:
            self.agent_log.log_stage_result("argumentation", result["output"], self.stage, False)

        return {"stage": self.stage, "output": result["output"], "auto": False}

    def _lg_data_confirmation(self, state: PipelineState) -> PipelineState:
        self.stage = state.get("stage", self.stage)
        self.name = self.stage

        text = self._get_effective_input(state)
        text = text.replace("/", "").replace("-", "").replace(".", "")

        if self.agent_log:
            self.agent_log.log_section("DATA_CONFIRMATION")
            self.agent_log.log("input", text)
            self.agent_log.log_messages("Histórico ANTES do run", self.agent.messages)

        t0 = time.perf_counter()
        result = self.agent.run_confirmation(text)
        elapsed = time.perf_counter() - t0

        if self.agent_log:
            self.agent_log.log("LLM tempo de execução", f"{elapsed:.3f}s")
            self.agent_log.log("LLM output", result["output"])
            self.agent_log.log("tool_calls", str(result.get("tools", [])))
            self.agent_log.log_agent_state("data_confirmation", self.agent.state)
            self.agent_log.log_messages("Histórico DEPOIS do run", self.agent.messages)

        # buffer do output gerado
        self._buffer_langfuse_trace(text, result["output"])

        if self.agent.state.is_authenticated:
            self.update_langfuse_auto(text, result["output"])
            self.stage = "formalization"

            self.agent = FormalizationAgent(
                streaming=self.streaming,
                prompt_vars=self.prompt_vars,
            )

            if self.agent_log:
                self.agent_log.log_stage_result("data_confirmation", result["output"], "formalization", True)
                self.agent_log.log_section("TRANSIÇÃO → FORMALIZATION")
                self.agent_log.log_system_prompt("formalization", self.agent.system_prompt)

            return {"stage": "formalization", "auto": True}

        if self.agent.state.end_conversation:
            self.stage = "DONE"
            if self.agent_log:
                self.agent_log.log_stage_result("data_confirmation", result["output"], "DONE", False)
            return {"stage": "DONE", "output": result["output"], "auto": False}

        if self.agent_log:
            self.agent_log.log_stage_result("data_confirmation", result["output"], self.stage, False)

        return {"stage": self.stage, "output": result["output"], "auto": False}

    def _lg_formalization(self, state: PipelineState) -> PipelineState:
        self.stage = state.get("stage", self.stage)
        self.name = self.stage

        text = self._get_effective_input(state)

        if self.agent_log:
            self.agent_log.log_section("FORMALIZATION")
            self.agent_log.log("input", text)
            self.agent_log.log_messages("Histórico ANTES do run", self.agent.messages)

        t0 = time.perf_counter()
        result = self.agent.run_formalization(text)
        elapsed = time.perf_counter() - t0

        if self.agent_log:
            self.agent_log.log("LLM tempo de execução", f"{elapsed:.3f}s")
            self.agent_log.log("LLM output", result["output"])
            self.agent_log.log("tool_calls", str(result.get("tools", [])))
            self.agent_log.log_agent_state("formalization", self.agent.state)
            self.agent_log.log_messages("Histórico DEPOIS do run", self.agent.messages)

        # buffer do output gerado
        self._buffer_langfuse_trace(text, result["output"])

        if self.agent.state.end_conversation:
            self.stage = "DONE"
            if self.agent_log:
                self.agent_log.log_stage_result("formalization", result["output"], "DONE (end_conversation)", False)
            return {"stage": "DONE", "output": result["output"], "auto": False}

        if self.agent.state.accepted:
            self.stage = "DONE"
            self.accepted_purchase = True
            if self.agent_log:
                self.agent_log.log_stage_result("formalization", result["output"], "DONE (compra aceita)", False)
            return {"stage": "DONE", "output": result["output"], "auto": False}

        if self.agent.state.declines >= 2:
            self.stage = "DONE"
            if self.agent_log:
                self.agent_log.log_stage_result("formalization", result["output"], f"DONE (declines={self.agent.state.declines})", False)
            return {"stage": "DONE", "output": result["output"], "auto": False}

        if self.agent_log:
            self.agent_log.log_stage_result("formalization", result["output"], self.stage, False)

        return {"stage": self.stage, "output": result["output"], "auto": False}

    # ====================================================================================
    # 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", ""

        if self.agent_log:
            self.agent_log.log_section(f"RUN (estágio: {self.stage})")
            self.agent_log.log("user_input", str(user_input))

        #if self.pending_output:
        #    self.set_interruption(was_interrupted=False)

        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)

        if self.agent_log:
            self.agent_log.log("run finalizado", f"stage={self.stage}  output={out.get('output', '')[:120]}...")

        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,
        }

        # Log de fim de serviço
        if self.agent_log:
            self.agent_log.log_section("END SERVICE")
            self.agent_log.log("accepted_purchase", str(self.accepted_purchase))
            self.agent_log.log("classified_conversation", str(self.classified_conversation))
            self.agent_log.log("summary", str(self.summary))
            self.agent_log.log("formatted_call", str(formatted_call))
            self.agent_log.log("success_purchase", str(success_purchase))

        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

    def get_formatted_metadata(self):
        if self.stage == "presentation":
            return {
                "stage": self.stage,
                "attacks": self.attacks,
                "presentation.name_validation_attempt": self.agent.state.name_validation_attempt,
                "presentation.eligibility": self.prompt_vars["elegibility"],
                "presentation.is_target_customer": self.agent.state.is_target_customer,
                "presentation.end_conversation": self.agent.state.end_conversation,
            }

        if self.stage == "argumentation":
            return {
                "stage": self.stage,
                "attacks": self.attacks,
                "argumentation.accepted": self.agent.state.accepted,
                "argumentation.accepted_count": self.agent.state.accepted_count,
                "argumentation.declines": self.agent.state.declines,
                "argumentation.end_conversation": self.agent.state.end_conversation,
                "argumentation.off_topic": self.agent.state.off_topic,
            }

        if self.stage == "data_confirmation":
            return {
                "stage": self.stage,
                "attacks": self.attacks,
                "data_confirmation.authenticated": self.agent.state.is_authenticated,
                "data_confirmation.validated_target_customer": self.agent.state.validated_target_customer,
                "data_confirmation.validated_cpf": self.agent.state.validated_cpf,
                "data_confirmation.validated_birth_date": self.agent.state.validated_birth_date,
                "data_confirmation.attempt_customer": self.agent.state.attempt_customer,
                "data_confirmation.attempt_cpf": self.agent.state.attempt_cpf,
                "data_confirmation.attempt_birth_date": self.agent.state.attempt_birth_date,
            }

        if self.stage == "formalization":
            return {
                "stage": self.stage,
                "attacks": self.attacks,
                "formalization.accepted": self.agent.state.accepted,
                "formalization.declines": self.agent.state.declines,
                "formalization.questions": self.agent.state.questions,
            }

        return {
            "stage": self.stage,
            "attacks": self.attacks,
        }
