# teste/teste.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

class TestState:
    def __init__(self):
        self.count = 0
        
class TestAgent(BaseAgent):
    def __init__(self, streaming=False, prompt_vars=None):
        
        state = TestState()

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

    def _build_tools(self):
        @tool
        def classify_customer_response(response: Literal["positive","negative", "other"]
        ):
            """Classifica a resposta do cliente"""

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

            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("test.txt").read_text(encoding="utf-8")
        return prompt_text

    def run(self, user_input):
        return super().run(user_input)

    def run_stream(self, user_input):
        return super().stream_run(user_input)