"""
Environment and Scope Management for the Alloy Programming Language

This module provides the Environment class which manages variable, tool, and agent
scoping and lookup in the Alloy interpreter.
"""

from typing import Any, Dict, List, Optional, Callable, Union


class RuntimeError(Exception):
    """Runtime error in Alloy execution"""
    pass


class Environment:
    """
    Environment for managing variables, tools, and agents in Alloy.
    
    The Environment class provides lexical scoping with support for:
    - Variable storage and lookup
    - Tool (function) definition and lookup
    - Agent definition and lookup
    - Parent environment chain for nested scopes
    """
    
    def __init__(self, parent: Optional['Environment'] = None):
        """
        Initialize a new environment.
        
        Args:
            parent: Parent environment for scope chain (None for global scope)
        """
        self.vars: Dict[str, Any] = {}
        self.tools: Dict[str, Callable] = {}
        self.agents: Dict[str, Dict[str, Any]] = {}  # Agent model names and tools
        self.commands: Dict[str, str] = {}  # Command definitions
        self.parent = parent

    def define(self, name: str, value: Any) -> None:
        """
        Define a variable in this environment.
        
        Args:
            name: Variable name
            value: Variable value
        """
        self.vars[name] = value

    def get(self, name: str) -> Any:
        """
        Get a variable value from this environment or parent chain.
        
        Args:
            name: Variable name
            
        Returns:
            Variable value
            
        Raises:
            RuntimeError: If variable is not defined
        """
        if name in self.vars:
            return self.vars[name]
        if self.parent:
            return self.parent.get(name)
        raise RuntimeError(f"Undefined variable '{name}'")

    def set(self, name: str, value: Any) -> None:
        """
        Set a variable value in the environment where it's defined.
        If not found in any environment, defines it in the current one.
        
        Args:
            name: Variable name
            value: Variable value
        """
        if name in self.vars:
            self.vars[name] = value
        elif self.parent and self.parent.has_var(name):
            self.parent.set(name, value)
        else:
            self.vars[name] = value

    def has_var(self, name: str) -> bool:
        """
        Check if a variable exists in this environment or parent chain.
        
        Args:
            name: Variable name
            
        Returns:
            True if variable exists, False otherwise
        """
        if name in self.vars:
            return True
        if self.parent:
            return self.parent.has_var(name)
        return False

    def define_tool(self, name: str, func: Callable) -> None:
        """
        Define a tool (function) in this environment.
        
        Args:
            name: Tool name
            func: Tool function
        """
        self.tools[name] = func

    def get_tool(self, name: str) -> Callable:
        """
        Get a tool function from this environment or parent chain.
        
        Args:
            name: Tool name
            
        Returns:
            Tool function
            
        Raises:
            RuntimeError: If tool is not defined
        """
        if name in self.tools:
            return self.tools[name]
        if self.parent:
            return self.parent.get_tool(name)
        raise RuntimeError(f"Undefined tool '{name}'")

    def has_tool(self, name: str) -> bool:
        """
        Check if a tool exists in this environment or parent chain.
        
        Args:
            name: Tool name
            
        Returns:
            True if tool exists, False otherwise
        """
        if name in self.tools:
            return True
        if self.parent:
            return self.parent.has_tool(name)
        return False

    def define_agent(self, name: str, agent_info: Union[str, Dict[str, Any]], tools: Optional[List[str]] = None) -> None:
        """
        Define an agent in this environment.
        
        Args:
            name: Agent name
            agent_info: Either model string (old format) or dict with model, tools, system_prompt (new format)
            tools: List of tool names the agent can access (only used with old format)
        """
        if isinstance(agent_info, str):
            # Old format: model string
            self.agents[name] = {
                'model': agent_info,
                'tools': tools or [],
                'system_prompt': "You are a helpful assistant."
            }
        else:
            # New format: dict with all info
            self.agents[name] = agent_info

    def get_agent(self, name: str) -> Dict[str, Any]:
        """
        Get agent definition from this environment or parent chain.
        
        Args:
            name: Agent name
            
        Returns:
            Agent definition dict with 'model' and 'tools' keys
            
        Raises:
            RuntimeError: If agent is not defined
        """
        if name in self.agents:
            return self.agents[name]
        if self.parent:
            return self.parent.get_agent(name)
        raise RuntimeError(f"Undefined agent '{name}'")

    def has_agent(self, name: str) -> bool:
        """
        Check if an agent exists in this environment or parent chain.
        
        Args:
            name: Agent name
            
        Returns:
            True if agent exists, False otherwise
        """
        if name in self.agents:
            return True
        if self.parent:
            return self.parent.has_agent(name)
        return False

    def define_command(self, name: str, description: str) -> None:
        """
        Define a command in this environment.
        
        Args:
            name: Command name
            description: Command description
        """
        self.commands[name] = description

    def get_command(self, name: str) -> str:
        """
        Get a command description from this environment or parent chain.
        
        Args:
            name: Command name
            
        Returns:
            Command description
            
        Raises:
            RuntimeError: If command is not defined
        """
        if name in self.commands:
            return self.commands[name]
        if self.parent:
            return self.parent.get_command(name)
        raise RuntimeError(f"Undefined command '{name}'")

    def has_command(self, name: str) -> bool:
        """
        Check if a command exists in this environment or parent chain.
        
        Args:
            name: Command name
            
        Returns:
            True if command exists, False otherwise
        """
        if name in self.commands:
            return True
        if self.parent:
            return self.parent.has_command(name)
        return False

    def get_all_vars(self) -> Dict[str, Any]:
        """
        Get all variables from this environment and parent chain.
        
        Returns:
            Dictionary of all variables (parent variables can be overridden)
        """
        all_vars = {}
        if self.parent:
            all_vars.update(self.parent.get_all_vars())
        all_vars.update(self.vars)
        return all_vars

    def get_all_tools(self) -> Dict[str, Callable]:
        """
        Get all tools from this environment and parent chain.
        
        Returns:
            Dictionary of all tools (parent tools can be overridden)
        """
        all_tools = {}
        if self.parent:
            all_tools.update(self.parent.get_all_tools())
        all_tools.update(self.tools)
        return all_tools

    def get_all_agents(self) -> Dict[str, Dict[str, Any]]:
        """
        Get all agents from this environment and parent chain.
        
        Returns:
            Dictionary of all agents (parent agents can be overridden)
        """
        all_agents = {}
        if self.parent:
            all_agents.update(self.parent.get_all_agents())
        all_agents.update(self.agents)
        return all_agents

    def get_all_commands(self) -> Dict[str, str]:
        """
        Get all commands from this environment and parent chain.
        
        Returns:
            Dictionary of all commands (parent commands can be overridden)
        """
        all_commands = {}
        if self.parent:
            all_commands.update(self.parent.get_all_commands())
        all_commands.update(self.commands)
        return all_commands

    def create_child(self) -> 'Environment':
        """
        Create a child environment with this environment as parent.
        
        Returns:
            New child environment
        """
        return Environment(parent=self)

    def __repr__(self) -> str:
        """String representation of environment for debugging"""
        return f"Environment(vars={list(self.vars.keys())}, tools={list(self.tools.keys())}, agents={list(self.agents.keys())}, commands={list(self.commands.keys())})"