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

class PresentationState:
    def __init__(self):
        self.name_validation_attempt = 0
        self.is_target_customer = False
        self.end_conversation = False

        self._saved_state = None
        self.save_state()

    def save_state(self):
        self._saved_state = copy.deepcopy(self.__dict__)
        self._saved_state.pop("_saved_state", None)

    def restore_state(self):
        if self._saved_state is None:
            raise RuntimeError("Nenhum estado foi salvo ainda.")

        self.__dict__.update(copy.deepcopy(self._saved_state))

    def print_state(self):
        print(f"name_validation_attempt: {self.name_validation_attempt}")
        print(f"is_target_customer: {self.is_target_customer}")
        print(f"end_conversation: {self.end_conversation}")

class PresentationAgent(BaseAgent):
    def __init__(self, streaming=False, prompt_vars=None):
        
        state = PresentationState()

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

    def _build_tools(self):
        @tool
        def classify_customer_response(response: Literal["confirmed_target_customer", 
                                                        "denied_target_customer",
                                                        "other"]
        ):
            """Classifica a resposta do interlocutor sobre ser ou não o cliente alvo.
confirmed_target_customer: confirmou ser o cliente alvo (ex: 'sou eu', 'sim', 'pode falar', disse o próprio nome)
denied_target_customer: negou ser o cliente alvo ou é terceiro (ex: 'não é ele', 'sou o filho', 'ele saiu')
other: resposta ambígua ou não relacionada à confirmação de identidade"""

            print(response)
            self.state.name_validation_attempt += 1

            if self.state.name_validation_attempt >= 4:
                return "Finalize a conversa com uma mensagem que contenha no final 'A TIM agradece sua atenção.'"
            
            if response == "confirmed_target_customer":
                self.state.is_target_customer = True
                return "###end_presentation###"
            
            return "Lide com o input do usuario"
        
        return [classify_customer_response]
    
    @lru_cache
    def _load_prompt(self) -> str:
        prompt_text = resources.files("agent.prompts").joinpath("presentation.txt").read_text(encoding="utf-8")
        general_rules = resources.files("agent.prompts").joinpath("general_rules.txt").read_text(encoding="utf-8")
        return prompt_text.replace("{general_rules}", general_rules)

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

        self.state.save_state()
        
        result = self.run(user_input)
                
        if "tim agradece sua atenção" in result['output'].lower():
            self.state.end_conversation = True
        
        return result