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

import copy

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

        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"declines: {self.declines}")
        print(f"accepted_count: {self.accepted_count}")
        print(f"accepted: {self.accepted}")
        print(f"end_conversation: {self.end_conversation}")
        print(f"off_topic: {self.off_topic}")

class ArgumentationAgent(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='argumentation'
        )
        self.state = state

    def _build_tools(self):
        @tool
        def detect_intention_purchase(intention: Literal["customer_declined_purchase", 
                                                         "customer_accepted_purchase",
                                                         "off_topic_from_sale",
                                                         "other"]):
            """Classifica a intenção do cliente em relação à oferta de troca de plano.
- customer_accepted_purchase: cliente demonstrou interesse claro em comprar o plano (ex: 'quero', 'pode ser', 'sim', 'topo')
- customer_declined_purchase: cliente recusou comprar o plano (ex: 'não quero', 'caro', 'não preciso', 'deixa esse plano')
- off_topic_from_sale: cliente falou sobre assuntos que não têm relação com a venda (ex: 'vou para a praia', 'vem almoçar')
- other: dúvidas, perguntas ou respostas ambíguas sobre o plano
            """  
            return _handle_intention(intention)
        
        def _handle_intention(intention):
            print(intention)
            #self.state.history.append(intention)

            if intention == "customer_accepted_purchase":
                self.state.accepted_count += 1
                if self.state.accepted_count >= 2:
                    self.state.accepted = True
                    return "###end_argumentation###"
                
                return "Peça a segunda confirmação"

            self.state.accepted_count = 0

            if intention == "customer_declined_purchase":
                self.state.declines += 1
                if self.state.declines > 2:
                    return f"{self.state.declines}º, finalize"
                return f"{self.state.declines}º rejeição"
            
            if intention == "off_topic_from_sale":
                self.state.off_topic += 1
                if self.state.off_topic > 3:
                    return "Finalize a conversa"
                
            return "Lide com o input do usuario"

        return [detect_intention_purchase]
    
    @lru_cache
    def _load_prompt(self) -> str:
        prompt_text = resources.files("agent.prompts").joinpath("argumentation.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 update_prompt_variable(self) -> None:
        pattern = r"(### Informações extras ###)(.*?)(### Fim das informações extras ###)"

        new_content = f"Recusas do cliente:{self.state.declines}"

        print(new_content)
        def replacer(match):
            return f"{match.group(1)}\n{new_content}\n{match.group(3)}"

        self.agent.steps[1].messages[0].prompt.template = re.sub(
            pattern,
            replacer,
            self.agent.steps[1].messages[0].prompt.template,
            flags=re.DOTALL
        )
        
    def run_argumentation(self, user_input):
        #self.update_prompt_variable()
        self.state.save_state()
        result = self.run(user_input)
        #print(self.messages)
        if ("tim agradece sua atenção" in result['output'].lower()):
            self.state.end_conversation = True

        return result