from typing import Dict, Optional

from pqai_agent.data_models.agent_configuration import AgentConfiguration
from pqai_agent.logging_service import logger

class AgentConfigManager:
    def __init__(self, session_maker):
        self.session_maker = session_maker
        self.agent_configs: Dict[int, AgentConfiguration] = {}
        self.refresh_configs()

    def refresh_configs(self):
        try:
            with self.session_maker() as session:
                data = session.query(AgentConfiguration).filter_by(is_delete=False).all()
                agent_configs = {}
                for config in data:
                    agent_configs[config.id] = config
            self.agent_configs = agent_configs
            logger.debug(f"Refreshed agent configurations: {list(self.agent_configs.keys())}")
        except Exception as e:
            logger.error(f"Failed to refresh agent configurations: {e}")

    def get_config(self, agent_id: int) -> Optional[AgentConfiguration]:
        return self.agent_configs.get(agent_id, None)