# robot_monitoring_tools/robot_monitoring_tools/agent_client.py
from dotenv import load_dotenv
from langchain_openai import AzureChatOpenAI
from langchain.agents import create_openai_functions_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from ..tools.battery_tool import CheckBatteryTool
from ..tools.transfer_date_tool import GetRobotLastTransferDateTool
import os

class AgentClient:
    def __init__(self, env_path: str = None, azure_api_key: str = None, azure_endpoint: str = None, tool_selection: list = None):
        """
        로봇 모니터링 에이전트 클라이언트 초기화
        
        Args:
            env_path (str, optional): .env 파일 경로
            azure_api_key (str): Azure API 키
            azure_endpoint (str): Azure 엔드포인트
            tool_selection (list, optional): 사용할 도구 목록, 비어있으면 모든 도구 사용
        """
        self.azure_api_key = azure_api_key  # Azure API 키 저장
        self.azure_endpoint = azure_endpoint  # Azure 엔드포인트 저장
        # Azure OpenAI 설정
        self.llm = self._setup_llm()
        
        # 도구 설정
        if tool_selection is None:
            self.tools = [CheckBatteryTool(), GetRobotLastTransferDateTool()]
        else : 
            self.tools = tool_selection
        # 프롬프트 설정
        self.prompt = self._setup_prompt()
        
        # 에이전트 설정
        self.agent = create_openai_functions_agent(self.llm, self.tools, self.prompt)
        self.agent_executor = AgentExecutor(agent=self.agent, tools=self.tools, verbose=True)
    def _setup_llm(self):
        """Azure OpenAI LLM 설정"""
        # api_key = os.getenv("AZURE_OPENAI_API_KEY")
        # azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
        
        if not self.azure_api_key or not self.azure_endpoint:
            raise ValueError("Azure OpenAI API 키와 엔드포인트가 필요합니다.")
            
        return AzureChatOpenAI(
            openai_api_version="2024-06-01",
            azure_deployment="gpt-4o",
            azure_endpoint=self.azure_endpoint,
            api_key=self.azure_api_key,
            temperature=0.7
        )

    def _setup_prompt(self):
        """프롬프트 템플릿 설정"""
        return ChatPromptTemplate.from_messages([
            ("system", "당신은 로봇 상태를 모니터링하는 도우미입니다."),
            ("user", "{input}"),
            ("assistant", "로봇의 상태를 확인하기 위해 적절한 도구를 사용하겠습니다."),
            MessagesPlaceholder(variable_name="agent_scratchpad")
        ])

    def ask(self, question: str) -> str:
        """
        에이전트에게 질문하고 응답 받기
        
        Args:
            question (str): 질문 내용
            
        Returns:
            str: 에이전트의 응답
        """
        response = self.agent_executor.invoke({"input": question})
        return response["output"]