# formalization/formalization.py
from functools import lru_cache
from importlib import resources
from langchain.tools import tool
from agent.base.base_stage import BaseAgent
from typing import Literal

class ArgumentationState:
    def __init__(self):
        self.declines = 0
        self.accepted = False
        self.questions = 0
        self.end_conversation = False

class FormalizationAgent(BaseAgent):
    def __init__(self, streaming=False, prompt_vars=None):
        state = ArgumentationState()

        tools = self._build_tools()
        super().__init__(tools=tools, 
                         streaming=streaming, 
                         prompt_vars=prompt_vars,
                         agent_name='formalization')
        self.state = state

    def _build_tools(self):
        @tool
        def detect_intention_purchase(intention: Literal["customer_declined_purchase", "customer_accepted_purchase", "customer_asked_question", "other"]):
            """Classifica a intenção do cliente durante a formalização da troca de plano.
- customer_accepted_purchase: cliente disse 'eu confirmo' aceitando a migração
- customer_declined_purchase: cliente recusou confirmar a migração (ex: 'não quero', 'mudei de ideia')
- customer_asked_question: cliente fez uma pergunta sobre o plano, benefícios ou cobrança
- other: resposta que não se encaixa nas categorias acima"""
            #print(">>>>",intention)
            if intention == "customer_declined_purchase":
                self.state.declines += 1
            if intention == "customer_accepted_purchase":
                self.state.accepted = True
            if intention == "customer_asked_question":
                self.state.questions += 1
                
            return intention
        @tool
        def purchase_cancellation(end: bool):
            """Encerra a chamada quando o cliente desistir claramente da compra. Chame com end=true."""
            #print(">>>>",end)
            self.state.end_conversation = end
            return end
        
        return [detect_intention_purchase, purchase_cancellation]

    @lru_cache
    def _load_prompt(self) -> str:
        prompt_text = resources.files("agent.prompts").joinpath("formalization.txt").read_text(encoding="utf-8")
        general_rules = resources.files("agent.prompts").joinpath("general_rules.txt").read_text(encoding="utf-8")
        knowledge_text = resources.files("agent.prompts").joinpath("knowledge_base.txt").read_text(encoding="utf-8")

        return prompt_text.replace("{general_rules}", general_rules).replace("{knowledge_base}", knowledge_text)

    def set_interruption_message(self, text: str):
        if self.messages:
            self.messages[-1].content = f"{text}"
            
    def run_formalization(self, user_input: str):
        result = self.run(user_input)
  
        return result
